jyapayne/langflow
1
0
Fork 0
Code Issues Pull requests Projects Releases 1 Packages Wiki Activity Actions

Merge remote-tracking branch 'origin/dev' into celery

Browse source
This commit is contained in:
Gabriel Luiz Freitas Almeida 2023-09-22 14:50:19 -03:00
commit b3febf25dd
31 changed files with 582 additions and 263 deletions
Download patch file Download diff file Expand all files Collapse all files

91
tests/conftest.py
View file

@ -2,13 +2,14 @@ from contextlib import contextmanager
import json
from pathlib import Path
from typing import AsyncGenerator, TYPE_CHECKING
from langflow.api.v1.flows import get_session
from langflow.graph.graph.base import Graph
from langflow.services.auth.utils import get_password_hash
from langflow.services.database.models.flow.flow import Flow, FlowCreate
from langflow.services.database.models.user.user import User, UserCreate
import orjson
from langflow.services.database.utils import session_getter
from langflow.services.getters import get_db_manager
import pytest
from fastapi.testclient import TestClient
from httpx import AsyncClient
@ -16,6 +17,9 @@ from sqlmodel import SQLModel, Session, create_engine
from sqlmodel.pool import StaticPool
from typer.testing import CliRunner
# we need to import tmpdir
import tempfile
if TYPE_CHECKING:
from langflow.services.database.manager import DatabaseService
@ -61,22 +65,6 @@ def session_fixture():
yield session
@pytest.fixture(name="client")
def client_fixture(session: Session, monkeypatch):
def get_session_override():
return session
monkeypatch.setenv("LANGFLOW_AUTO_LOGIN", False)
from langflow.main import create_app
app = create_app()
app.dependency_overrides[get_session] = get_session_override
with TestClient(app) as client:
yield client
app.dependency_overrides.clear()
class Config:
broker_url = "redis://localhost:6379/0"
result_backend = "redis://localhost:6379/0"
@ -182,27 +170,33 @@ def json_vector_store():
return f.read()
# @contextmanager
# def session_getter():
# try:
# session = Session(engine)
# yield session
# except Exception as e:
# print("Session rollback because of exception:", e)
# session.rollback()
# raise
# finally:
# session.close()
@pytest.fixture(name="client", autouse=True)
def client_fixture(session: Session, monkeypatch):
# Set the database url to a test database
db_dir = tempfile.mkdtemp()
db_path = Path(db_dir) / "test.db"
monkeypatch.setenv("LANGFLOW_DATABASE_URL", f"sqlite:///{db_path}")
monkeypatch.setenv("LANGFLOW_AUTO_LOGIN", False)
def get_session_override():
return session
from langflow.main import create_app
app = create_app()
# app.dependency_overrides[get_session] = get_session_override
with TestClient(app) as client:
yield client
# app.dependency_overrides.clear()
monkeypatch.undo()
# clear the temp db
db_path.unlink()
# create a fixture for session_getter above
@pytest.fixture(name="session_getter")
def session_getter_fixture(client):
engine = create_engine(
"sqlite://", connect_args={"check_same_thread": False}, poolclass=StaticPool
)
SQLModel.metadata.create_all(engine)
@contextmanager
def blank_session_getter(db_service: "DatabaseService"):
with Session(db_service.engine) as session:
@ -228,17 +222,18 @@ def test_user(client):
@pytest.fixture(scope="function")
def active_user(client, session):
user = User(
username="activeuser",
password=get_password_hash(
"testpassword"
), # Assuming password needs to be hashed
is_active=True,
is_superuser=False,
)
session.add(user)
session.commit()
def active_user(client):
db_manager = get_db_manager()
with session_getter(db_manager) as session:
user = User(
username="activeuser",
password=get_password_hash("testpassword"),
is_active=True,
is_superuser=False,
)
session.add(user)
session.commit()
session.refresh(user)
return user
@ -253,7 +248,7 @@ def logged_in_headers(client, active_user):
@pytest.fixture
def flow(client, json_flow: str, session, active_user):
def flow(client, json_flow: str, active_user):
from langflow.services.database.models.flow.flow import FlowCreate
loaded_json = json.loads(json_flow)
@ -261,8 +256,10 @@ def flow(client, json_flow: str, session, active_user):
name="test_flow", data=loaded_json.get("data"), user_id=active_user.id
)
flow = Flow(**flow_data.dict())
session.add(flow)
session.commit()
with session_getter(get_db_manager()) as session:
session.add(flow)
session.commit()
session.refresh(flow)
return flow

4
tests/test_cli.py
View file

@ -3,7 +3,7 @@ from tempfile import tempdir
from langflow.__main__ import app
import pytest
from langflow.services import utils
from langflow.services import getters
@pytest.fixture(scope="module")
@ -26,7 +26,7 @@ def test_components_path(runner, client, default_settings):
["run", "--components-path", str(temp_dir), *default_settings],
)
assert result.exit_code == 0, result.stdout
settings_service = utils.get_settings_service()
settings_service = getters.get_settings_service()
assert str(temp_dir) in settings_service.settings.COMPONENTS_PATH

8
tests/test_custom_component.py
View file

@ -518,13 +518,13 @@ def db(app):
app.db.drop_all()
def test_list_flows_return_type(component, session_getter):
flows = component.list_flows(get_session=session_getter)
def test_list_flows_return_type(component):
flows = component.list_flows()
assert isinstance(flows, list)
def test_list_flows_flow_objects(component, session_getter):
flows = component.list_flows(get_session=session_getter)
def test_list_flows_flow_objects(component):
flows = component.list_flows()
assert all(isinstance(flow, Flow) for flow in flows)

18
tests/test_database.py
View file

@ -1,4 +1,6 @@
from langflow.services.database.models.base import orjson_dumps
from langflow.services.database.utils import session_getter
from langflow.services.getters import get_db_manager
import orjson
import pytest
@ -194,18 +196,20 @@ def test_download_file(
FlowCreate(name="Flow 2", description="description", data=data),
]
)
for flow in flow_list.flows:
flow.user_id = active_user.id
db_flow = Flow.from_orm(flow)
session.add(db_flow)
session.commit()
db_manager = get_db_manager()
with session_getter(db_manager) as session:
for flow in flow_list.flows:
flow.user_id = active_user.id
db_flow = Flow.from_orm(flow)
session.add(db_flow)
session.commit()
# Make request to endpoint
response = client.get("api/v1/flows/download/", headers=logged_in_headers)
# Check response status code
assert response.status_code == 200
assert response.status_code == 200, response.json()
# Check response data
response_data = response.json()["flows"]
assert len(response_data) == 2
assert len(response_data) == 2, response_data
assert response_data[0]["name"] == "Flow 1"
assert response_data[0]["description"] == "description"
assert response_data[0]["data"] == data

15
tests/test_endpoints.py
View file

@ -2,7 +2,9 @@ import uuid
from langflow.processing.process import Result
from langflow.services.auth.utils import get_password_hash
from langflow.services.database.models.api_key.api_key import ApiKey
from langflow.services.utils import get_settings_service
from langflow.services.getters import get_settings_service
from langflow.services.database.utils import session_getter
from langflow.services.getters import get_db_manager
import pytest
from fastapi.testclient import TestClient
from langflow.interface.tools.constants import CUSTOM_TOOLS
@ -117,7 +119,7 @@ PROMPT_REQUEST = {
@pytest.fixture
def created_api_key(session, active_user):
def created_api_key(active_user):
hashed = get_password_hash("random_key")
api_key = ApiKey(
name="test_api_key",
@ -125,10 +127,11 @@ def created_api_key(session, active_user):
api_key="random_key",
hashed_api_key=hashed,
)
session.add(api_key)
session.commit()
session.refresh(api_key)
db_manager = get_db_manager()
with session_getter(db_manager) as session:
session.add(api_key)
session.commit()
session.refresh(api_key)
return api_key

9
tests/test_login.py
View file

@ -1,3 +1,5 @@
from langflow.services.database.utils import session_getter
from langflow.services.getters import get_db_manager
import pytest
from langflow.services.database.models.user import User
from langflow.services.auth.utils import get_password_hash
@ -15,10 +17,11 @@ def test_user():
)
def test_login_successful(client, test_user, session):
def test_login_successful(client, test_user):
# Adding the test user to the database
session.add(test_user)
session.commit()
with session_getter(get_db_manager()) as session:
session.add(test_user)
session.commit()
response = client.post(
"api/v1/login", data={"username": "testuser", "password": "testpassword"}

2
tests/test_prompts_template.py
View file

@ -1,5 +1,5 @@
from fastapi.testclient import TestClient
from langflow.services.utils import get_settings_service
from langflow.services.getters import get_settings_service
def test_prompts_settings(client: TestClient, logged_in_headers):

140
tests/test_setup_superuser.py Normal file
View file

@ -0,0 +1,140 @@
from unittest.mock import patch, Mock, MagicMock, call
from langflow.services.database.models.user.user import User
from langflow.services.settings.constants import (
DEFAULT_SUPERUSER,
DEFAULT_SUPERUSER_PASSWORD,
)
from langflow.services.utils import setup_superuser, teardown_superuser
@patch("langflow.services.utils.get_settings_manager")
@patch("langflow.services.utils.create_super_user")
@patch("langflow.services.utils.get_session")
def test_setup_superuser(
mock_get_session, mock_create_super_user, mock_get_settings_manager
):
# Test when AUTO_LOGIN is True
calls = []
mock_settings_manager = Mock()
mock_settings_manager.auth_settings.AUTO_LOGIN = True
mock_settings_manager.auth_settings.SUPERUSER = DEFAULT_SUPERUSER
mock_settings_manager.auth_settings.SUPERUSER_PASSWORD = DEFAULT_SUPERUSER_PASSWORD
mock_get_settings_manager.return_value = mock_settings_manager
mock_session = Mock()
mock_session.query.return_value.filter.return_value.first.return_value = (
mock_session
)
# return value of get_session is a generator
mock_get_session.return_value = iter([mock_session, mock_session, mock_session])
setup_superuser()
mock_session.query.assert_called_once_with(User)
actual_expr = mock_session.query.return_value.filter.call_args[0][0]
expected_expr = User.username == DEFAULT_SUPERUSER
assert str(actual_expr) == str(expected_expr)
create_call = call(
db=mock_session, username=DEFAULT_SUPERUSER, password=DEFAULT_SUPERUSER_PASSWORD
)
calls.append(create_call)
mock_create_super_user.assert_has_calls(calls)
assert 1 == mock_create_super_user.call_count
def reset_mock_credentials():
mock_settings_manager.auth_settings.SUPERUSER = DEFAULT_SUPERUSER
mock_settings_manager.auth_settings.SUPERUSER_PASSWORD = (
DEFAULT_SUPERUSER_PASSWORD
)
ADMIN_USER_NAME = "admin_user"
# Test when username and password are default
mock_settings_manager.auth_settings = Mock()
mock_settings_manager.auth_settings.AUTO_LOGIN = False
mock_settings_manager.auth_settings.SUPERUSER = ADMIN_USER_NAME
mock_settings_manager.auth_settings.SUPERUSER_PASSWORD = "password"
mock_settings_manager.auth_settings.reset_credentials = Mock(
side_effect=reset_mock_credentials
)
mock_get_settings_manager.return_value = mock_settings_manager
setup_superuser()
mock_session.query.assert_called_with(User)
actual_expr = mock_session.query.return_value.filter.call_args[0][0]
expected_expr = User.username == ADMIN_USER_NAME
assert str(actual_expr) == str(expected_expr)
create_call = call(db=mock_session, username=ADMIN_USER_NAME, password="password")
calls.append(create_call)
mock_create_super_user.assert_has_calls(calls)
assert 2 == mock_create_super_user.call_count
# Test that superuser credentials are reset
mock_settings_manager.auth_settings.reset_credentials.assert_called_once()
assert mock_settings_manager.auth_settings.SUPERUSER != ADMIN_USER_NAME
assert mock_settings_manager.auth_settings.SUPERUSER_PASSWORD != "password"
# Test when superuser already exists
mock_settings_manager.auth_settings.AUTO_LOGIN = False
mock_settings_manager.auth_settings.SUPERUSER = ADMIN_USER_NAME
mock_settings_manager.auth_settings.SUPERUSER_PASSWORD = "password"
mock_user = Mock()
mock_user.is_superuser = True
mock_session.query.return_value.filter.return_value.first.return_value = mock_user
setup_superuser()
mock_session.query.assert_called_with(User)
actual_expr = mock_session.query.return_value.filter.call_args[0][0]
expected_expr = User.username == ADMIN_USER_NAME
assert str(actual_expr) == str(expected_expr)
@patch("langflow.services.utils.get_settings_manager")
@patch("langflow.services.utils.get_session")
def test_teardown_superuser_default_superuser(
mock_get_session, mock_get_settings_manager
):
mock_settings_manager = MagicMock()
mock_settings_manager.auth_settings.AUTO_LOGIN = True
mock_settings_manager.auth_settings.SUPERUSER = DEFAULT_SUPERUSER
mock_settings_manager.auth_settings.SUPERUSER_PASSWORD = DEFAULT_SUPERUSER_PASSWORD
mock_get_settings_manager.return_value = mock_settings_manager
mock_session = MagicMock()
mock_user = MagicMock()
mock_user.is_superuser = True
mock_session.query.return_value.filter.return_value.first.return_value = mock_user
mock_get_session.return_value = iter([mock_session])
teardown_superuser()
mock_session.query.assert_called_once_with(User)
actual_expr = mock_session.query.return_value.filter.call_args[0][0]
expected_expr = User.username == DEFAULT_SUPERUSER
assert str(actual_expr) == str(expected_expr)
mock_session.delete.assert_called_once_with(mock_user)
mock_session.commit.assert_called_once()
@patch("langflow.services.utils.get_settings_manager")
@patch("langflow.services.utils.get_session")
def test_teardown_superuser_no_default_superuser(
mock_get_session, mock_get_settings_manager
):
ADMIN_USER_NAME = "admin_user"
mock_settings_manager = MagicMock()
mock_settings_manager.auth_settings.AUTO_LOGIN = False
mock_settings_manager.auth_settings.SUPERUSER = ADMIN_USER_NAME
mock_settings_manager.auth_settings.SUPERUSER_PASSWORD = "password"
mock_get_settings_manager.return_value = mock_settings_manager
mock_session = MagicMock()
mock_user = MagicMock()
mock_user.is_superuser = False
mock_session.query.return_value.filter.return_value.first.return_value = mock_user
mock_get_session.return_value = [mock_session]
teardown_superuser()
mock_session.query.assert_not_called()
mock_session.delete.assert_not_called()
mock_session.commit.assert_not_called()

86
tests/test_user.py
View file

@ -2,20 +2,22 @@ from datetime import datetime
from langflow.services.auth.utils import create_super_user, get_password_hash
from langflow.services.database.models.user.user import User
from langflow.services.utils import get_settings_service
from langflow.services.database.utils import session_getter
from langflow.services.getters import get_db_service, get_settings_service
import pytest
from langflow.services.database.models.user import UserUpdate
@pytest.fixture
def super_user(client, session):
settings_service = get_settings_service()
auth_settings = settings_service.auth_settings
return create_super_user(
db=session,
username=auth_settings.FIRST_SUPERUSER,
password=auth_settings.FIRST_SUPERUSER_PASSWORD,
)
def super_user(client):
settings_manager = get_settings_service()
auth_settings = settings_manager.auth_settings
with session_getter(get_db_service()) as session:
return create_super_user(
db=session,
username=auth_settings.SUPERUSER,
password=auth_settings.SUPERUSER_PASSWORD,
)
@pytest.fixture
@ -23,8 +25,8 @@ def super_user_headers(client, super_user):
settings_service = get_settings_service()
auth_settings = settings_service.auth_settings
login_data = {
"username": auth_settings.FIRST_SUPERUSER,
"password": auth_settings.FIRST_SUPERUSER_PASSWORD,
"username": auth_settings.SUPERUSER,
"password": auth_settings.SUPERUSER_PASSWORD,
}
response = client.post("/api/v1/login", data=login_data)
assert response.status_code == 200
@ -34,29 +36,34 @@ def super_user_headers(client, super_user):
@pytest.fixture
def deactivated_user(session):
user = User(
username="deactivateduser",
password=get_password_hash("testpassword"),
is_active=False,
is_superuser=False,
last_login_at=datetime.now(),
)
session.add(user)
session.commit()
def deactivated_user():
with session_getter(get_db_manager()) as session:
user = User(
username="deactivateduser",
password=get_password_hash("testpassword"),
is_active=False,
is_superuser=False,
last_login_at=datetime.now(),
)
session.add(user)
session.commit()
session.refresh(user)
return user
def test_user_waiting_for_approval(client, session):
def test_user_waiting_for_approval(
client,
):
# Create a user that is not active and has never logged in
user = User(
username="waitingforapproval",
password=get_password_hash("testpassword"),
is_active=False,
last_login_at=None,
)
session.add(user)
session.commit()
with session_getter(get_db_manager()) as session:
user = User(
username="waitingforapproval",
password=get_password_hash("testpassword"),
is_active=False,
last_login_at=None,
)
session.add(user)
session.commit()
login_data = {"username": "waitingforapproval", "password": "testpassword"}
response = client.post("/api/v1/login", data=login_data)
@ -106,16 +113,17 @@ def test_data_consistency_after_delete(client, test_user, super_user_headers):
assert all(user["id"] != user_id for user in response.json()["users"])
def test_inactive_user(client, session):
def test_inactive_user(client):
# Create a user that is not active and has a last_login_at value
user = User(
username="inactiveuser",
password=get_password_hash("testpassword"),
is_active=False,
last_login_at="2023-01-01T00:00:00", # Set to a valid datetime string
)
session.add(user)
session.commit()
with session_getter(get_db_manager()) as session:
user = User(
username="inactiveuser",
password=get_password_hash("testpassword"),
is_active=False,
last_login_at="2023-01-01T00:00:00", # Set to a valid datetime string
)
session.add(user)
session.commit()
login_data = {"username": "inactiveuser", "password": "testpassword"}
response = client.post("/api/v1/login", data=login_data)

2
tests/test_vectorstore_template.py
View file

@ -1,5 +1,5 @@
from fastapi.testclient import TestClient
from langflow.services.utils import get_settings_service
from langflow.services.getters import get_settings_service
# check that all agents are in settings.agents