tests: update env.py and conftest so tests use unique databases (#3654)
* Refactor Alembic environment script to simplify database connection logic and remove unused imports * Refactor test client fixture to use TemporaryDirectory and UUID for database path * Add `reload_engine` method to reinitialize the database engine * Add cleanup steps to test fixtures to ensure proper resource management - Added cleanup steps to various pytest fixtures to ensure temporary directories, database entries, and other resources are properly cleaned up after tests. - Modified fixtures to use `yield` for better resource management. - Ensured database connections are closed and tables are dropped after tests. - Improved temporary directory handling with context managers. * Add unit test to verify return type of list_flows method in custom component * Refactor tests in `test_custom_component.py` to remove unused imports and fixtures, and update existing fixtures for consistency. * Add debug checks for user existence in `test_user_waiting_for_approval` test * Fix import order and add flow_id validation in transaction handling - Corrected the import order in `utils.py`. - Added validation for `flow_id` in the transaction handling logic to ensure it is set correctly. * Add function to delete transactions by flow ID in CRUD module * Add cleanup for transactions and vertex builds in test teardown - Introduced `_delete_transactions_and_vertex_builds` function to remove transactions and vertex builds associated with a user. - Updated `user` fixture to call the new cleanup function before deleting the user. * Refactor flow_id assignment logic in `utils.py` to improve readability and correctness * [autofix.ci] apply automated fixes * Refactor test to use pytest's tmp_path fixture for temporary directory creation * Convert `test_user_waiting_for_approval` to an async test function --------- Co-authored-by: autofix-ci[bot] <114827586+autofix-ci[bot]@users.noreply.github.com>
This commit is contained in:
Gabriel Luiz Freitas Almeida
GitHub
parent
4e36dcc2ad
commit
2adda780c9
9 changed files with 179 additions and 149 deletions
|
|
@ -8,6 +8,7 @@ from collections.abc import AsyncGenerator
|
|||
from contextlib import contextmanager, suppress
|
||||
from pathlib import Path
|
||||
from typing import TYPE_CHECKING
|
||||
from uuid import UUID
|
||||
|
||||
import orjson
|
||||
import pytest
|
||||
|
|
@ -29,7 +30,9 @@ from langflow.services.auth.utils import get_password_hash
|
|||
from langflow.services.database.models.api_key.model import ApiKey
|
||||
from langflow.services.database.models.flow.model import Flow, FlowCreate
|
||||
from langflow.services.database.models.folder.model import Folder
|
||||
from langflow.services.database.models.user.model import User, UserCreate
|
||||
from langflow.services.database.models.transactions.model import TransactionTable
|
||||
from langflow.services.database.models.user.model import User, UserCreate, UserRead
|
||||
from langflow.services.database.models.vertex_builds.crud import delete_vertex_builds_by_flow_id
|
||||
from langflow.services.database.utils import session_getter
|
||||
from langflow.services.deps import get_db_service
|
||||
|
||||
|
|
@ -82,6 +85,23 @@ def get_text():
|
|||
assert path.exists(), f"File {path} does not exist. Available files: {list(data_path.iterdir())}"
|
||||
|
||||
|
||||
def delete_transactions_by_flow_id(db: Session, flow_id: UUID):
|
||||
stmt = select(TransactionTable).where(TransactionTable.flow_id == flow_id)
|
||||
transactions = db.exec(stmt)
|
||||
for transaction in transactions:
|
||||
db.delete(transaction)
|
||||
db.commit()
|
||||
|
||||
|
||||
def _delete_transactions_and_vertex_builds(session, user: User):
|
||||
flow_ids = [flow.id for flow in user.flows]
|
||||
for flow_id in flow_ids:
|
||||
if not flow_id:
|
||||
continue
|
||||
delete_vertex_builds_by_flow_id(session, flow_id)
|
||||
delete_transactions_by_flow_id(session, flow_id)
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def caplog(caplog: LogCaptureFixture):
|
||||
handler_id = logger.add(
|
||||
|
|
@ -110,6 +130,7 @@ def session_fixture():
|
|||
SQLModel.metadata.create_all(engine)
|
||||
with Session(engine) as session:
|
||||
yield session
|
||||
SQLModel.metadata.drop_all(engine) # Add this line to clean up tables
|
||||
|
||||
|
||||
class Config:
|
||||
|
|
@ -119,8 +140,8 @@ class Config:
|
|||
|
||||
@pytest.fixture(name="load_flows_dir")
|
||||
def load_flows_dir():
|
||||
tempdir = tempfile.TemporaryDirectory()
|
||||
yield tempdir.name
|
||||
with tempfile.TemporaryDirectory() as tempdir:
|
||||
yield tempdir
|
||||
|
||||
|
||||
@pytest.fixture(name="distributed_env")
|
||||
|
|
@ -143,23 +164,26 @@ def distributed_client_fixture(session: Session, monkeypatch, distributed_env):
|
|||
from langflow.core import celery_app
|
||||
|
||||
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")
|
||||
# monkeypatch langflow.services.task.manager.USE_CELERY to True
|
||||
# monkeypatch.setattr(manager, "USE_CELERY", True)
|
||||
monkeypatch.setattr(celery_app, "celery_app", celery_app.make_celery("langflow", Config))
|
||||
try:
|
||||
db_path = Path(db_dir) / "test.db"
|
||||
monkeypatch.setenv("LANGFLOW_DATABASE_URL", f"sqlite:///{db_path}")
|
||||
monkeypatch.setenv("LANGFLOW_AUTO_LOGIN", "false")
|
||||
# monkeypatch langflow.services.task.manager.USE_CELERY to True
|
||||
# monkeypatch.setattr(manager, "USE_CELERY", True)
|
||||
monkeypatch.setattr(celery_app, "celery_app", celery_app.make_celery("langflow", Config))
|
||||
|
||||
# def get_session_override():
|
||||
# return session
|
||||
# def get_session_override():
|
||||
# return session
|
||||
|
||||
from langflow.main import create_app
|
||||
from langflow.main import create_app
|
||||
|
||||
app = create_app()
|
||||
app = create_app()
|
||||
|
||||
# app.dependency_overrides[get_session] = get_session_override
|
||||
with TestClient(app) as client:
|
||||
yield client
|
||||
# app.dependency_overrides[get_session] = get_session_override
|
||||
with TestClient(app) as client:
|
||||
yield client
|
||||
finally:
|
||||
shutil.rmtree(db_dir) # Clean up the temporary directory
|
||||
app.dependency_overrides.clear()
|
||||
monkeypatch.undo()
|
||||
|
||||
|
|
@ -279,7 +303,9 @@ async def client_fixture(session: Session, monkeypatch, request, load_flows_dir)
|
|||
from langflow.main import create_app
|
||||
|
||||
app = create_app()
|
||||
|
||||
db_service = get_db_service()
|
||||
db_service.database_url = f"sqlite:///{db_path}"
|
||||
db_service.reload_engine()
|
||||
# app.dependency_overrides[get_session] = get_session_override
|
||||
async with LifespanManager(app, startup_timeout=None, shutdown_timeout=None) as manager:
|
||||
async with AsyncClient(transport=ASGITransport(app=manager.app), base_url="http://testserver/") as client:
|
||||
|
|
@ -304,7 +330,7 @@ def session_getter_fixture(client):
|
|||
|
||||
@pytest.fixture
|
||||
def runner():
|
||||
return CliRunner()
|
||||
yield CliRunner()
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
|
|
@ -315,26 +341,38 @@ async def test_user(client):
|
|||
)
|
||||
response = await client.post("api/v1/users/", json=user_data.model_dump())
|
||||
assert response.status_code == 201
|
||||
return response.json()
|
||||
user = response.json()
|
||||
yield user
|
||||
# Clean up
|
||||
await client.delete(f"/api/v1/users/{user['id']}")
|
||||
|
||||
|
||||
@pytest.fixture(scope="function")
|
||||
def active_user(client):
|
||||
db_manager = get_db_service()
|
||||
with session_getter(db_manager) as session:
|
||||
with db_manager.with_session() as session:
|
||||
user = User(
|
||||
username="activeuser",
|
||||
password=get_password_hash("testpassword"),
|
||||
is_active=True,
|
||||
is_superuser=False,
|
||||
)
|
||||
# check if user exists
|
||||
if active_user := session.exec(select(User).where(User.username == user.username)).first():
|
||||
return active_user
|
||||
session.add(user)
|
||||
user = active_user
|
||||
else:
|
||||
session.add(user)
|
||||
session.commit()
|
||||
session.refresh(user)
|
||||
user = UserRead.model_validate(user, from_attributes=True)
|
||||
yield user
|
||||
# Clean up
|
||||
# Now cleanup transactions, vertex_build
|
||||
with db_manager.with_session() as session:
|
||||
user = session.get(User, user.id)
|
||||
_delete_transactions_and_vertex_builds(session, user)
|
||||
session.delete(user)
|
||||
|
||||
session.commit()
|
||||
session.refresh(user)
|
||||
return user
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
|
|
@ -344,7 +382,7 @@ async def logged_in_headers(client, active_user):
|
|||
assert response.status_code == 200
|
||||
tokens = response.json()
|
||||
a_token = tokens["access_token"]
|
||||
return {"Authorization": f"Bearer {a_token}"}
|
||||
yield {"Authorization": f"Bearer {a_token}"}
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
|
|
@ -359,20 +397,22 @@ def flow(client, json_flow: str, active_user):
|
|||
session.add(flow)
|
||||
session.commit()
|
||||
session.refresh(flow)
|
||||
|
||||
return flow
|
||||
yield flow
|
||||
# Clean up
|
||||
session.delete(flow)
|
||||
session.commit()
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def json_chat_input():
|
||||
with open(pytest.CHAT_INPUT) as f:
|
||||
return f.read()
|
||||
yield f.read()
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def json_two_outputs():
|
||||
with open(pytest.TWO_OUTPUTS) as f:
|
||||
return f.read()
|
||||
yield f.read()
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
|
|
@ -385,7 +425,7 @@ async def added_flow_webhook_test(client, json_webhook_test, logged_in_headers):
|
|||
assert response.json()["name"] == flow.name
|
||||
assert response.json()["data"] == flow.data
|
||||
yield response.json()
|
||||
client.delete(f"api/v1/flows/{response.json()['id']}", headers=logged_in_headers)
|
||||
await client.delete(f"api/v1/flows/{response.json()['id']}", headers=logged_in_headers)
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
|
|
@ -398,7 +438,7 @@ async def added_flow_chat_input(client, json_chat_input, logged_in_headers):
|
|||
assert response.json()["name"] == flow.name
|
||||
assert response.json()["data"] == flow.data
|
||||
yield response.json()
|
||||
client.delete(f"api/v1/flows/{response.json()['id']}", headers=logged_in_headers)
|
||||
await client.delete(f"api/v1/flows/{response.json()['id']}", headers=logged_in_headers)
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
|
|
@ -411,7 +451,7 @@ async def added_flow_two_outputs(client, json_two_outputs, logged_in_headers):
|
|||
assert response.json()["name"] == flow.name
|
||||
assert response.json()["data"] == flow.data
|
||||
yield response.json()
|
||||
client.delete(f"api/v1/flows/{response.json()['id']}", headers=logged_in_headers)
|
||||
await client.delete(f"api/v1/flows/{response.json()['id']}", headers=logged_in_headers)
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
|
|
@ -424,7 +464,7 @@ async def added_vector_store(client, json_vector_store, logged_in_headers):
|
|||
assert response.json()["name"] == vector_store.name
|
||||
assert response.json()["data"] == vector_store.data
|
||||
yield response.json()
|
||||
client.delete(f"api/v1/flows/{response.json()['id']}", headers=logged_in_headers)
|
||||
await client.delete(f"api/v1/flows/{response.json()['id']}", headers=logged_in_headers)
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
|
|
@ -439,7 +479,7 @@ async def added_webhook_test(client, json_webhook_test, logged_in_headers):
|
|||
assert response.json()["name"] == webhook_test.name
|
||||
assert response.json()["data"] == webhook_test.data
|
||||
yield response.json()
|
||||
client.delete(f"api/v1/flows/{response.json()['id']}", headers=logged_in_headers)
|
||||
await client.delete(f"api/v1/flows/{response.json()['id']}", headers=logged_in_headers)
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
|
|
@ -451,7 +491,7 @@ async def flow_component(client: TestClient, logged_in_headers):
|
|||
response = await client.post("api/v1/flows/", json=flow.model_dump(), headers=logged_in_headers)
|
||||
assert response.status_code == 201
|
||||
yield response.json()
|
||||
client.delete(f"api/v1/flows/{response.json()['id']}", headers=logged_in_headers)
|
||||
await client.delete(f"api/v1/flows/{response.json()['id']}", headers=logged_in_headers)
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
|
|
@ -466,11 +506,15 @@ def created_api_key(active_user):
|
|||
db_manager = get_db_service()
|
||||
with session_getter(db_manager) as session:
|
||||
if existing_api_key := session.exec(select(ApiKey).where(ApiKey.api_key == api_key.api_key)).first():
|
||||
return existing_api_key
|
||||
yield existing_api_key
|
||||
return
|
||||
session.add(api_key)
|
||||
session.commit()
|
||||
session.refresh(api_key)
|
||||
return api_key
|
||||
yield api_key
|
||||
# Clean up
|
||||
session.delete(api_key)
|
||||
session.commit()
|
||||
|
||||
|
||||
@pytest.fixture(name="simple_api_test")
|
||||
|
|
@ -483,7 +527,7 @@ async def get_simple_api_test(client, logged_in_headers, json_simple_api_test):
|
|||
response = await client.post("api/v1/flows/", json=flow.model_dump(), headers=logged_in_headers)
|
||||
assert response.status_code == 201
|
||||
yield response.json()
|
||||
client.delete(f"api/v1/flows/{response.json()['id']}", headers=logged_in_headers)
|
||||
await client.delete(f"api/v1/flows/{response.json()['id']}", headers=logged_in_headers)
|
||||
|
||||
|
||||
@pytest.fixture(name="starter_project")
|
||||
|
|
@ -511,4 +555,7 @@ def get_starter_project(active_user):
|
|||
session.commit()
|
||||
session.refresh(new_flow)
|
||||
new_flow_dict = new_flow.model_dump()
|
||||
return new_flow_dict
|
||||
yield new_flow_dict
|
||||
# Clean up
|
||||
session.delete(new_flow)
|
||||
session.commit()
|
||||
|
|
|
|||
Loading…
Add table
Add a link
Reference in a new issue