From 2da18722d42d193afbb650133b2a856dd3d55069 Mon Sep 17 00:00:00 2001 From: nishika26 Date: Thu, 10 Jul 2025 14:55:01 +0530 Subject: [PATCH 1/2] transactioal db --- backend/app/tests/conftest.py | 150 +++++++++++++++++++++------------- 1 file changed, 91 insertions(+), 59 deletions(-) diff --git a/backend/app/tests/conftest.py b/backend/app/tests/conftest.py index 9f6fe76c..f6ada599 100644 --- a/backend/app/tests/conftest.py +++ b/backend/app/tests/conftest.py @@ -1,65 +1,97 @@ -from collections.abc import Generator - -import pytest +from uuid import uuid4 +from datetime import datetime, timezone from fastapi.testclient import TestClient -from sqlmodel import Session, delete - +from sqlmodel import Session from app.core.config import settings -from app.core.db import engine, init_db +from app.models import Collection from app.main import app -from app.models import ( - APIKey, - Assistant, - Organization, - Project, - ProjectUser, - User, - OpenAI_Thread, - Credential, - Collection, -) -from app.tests.utils.user import authentication_token_from_email -from app.tests.utils.utils import get_superuser_token_headers -from app.seed_data.seed_data import seed_database - - -@pytest.fixture(scope="session", autouse=True) -def db() -> Generator[Session, None, None]: - with Session(engine) as session: - init_db(session) - yield session - # Delete data in reverse dependency order - session.execute(delete(ProjectUser)) # Many-to-many relationship - session.execute(delete(Assistant)) - session.execute(delete(Credential)) - session.execute(delete(Project)) - session.execute(delete(Organization)) - session.execute(delete(APIKey)) - session.execute(delete(User)) - session.execute(delete(OpenAI_Thread)) - session.execute(delete(Collection)) - session.commit() - - -@pytest.fixture(scope="module") -def client() -> Generator[TestClient, None, None]: - with TestClient(app) as c: - yield c - - -@pytest.fixture(scope="module") -def superuser_token_headers(client: TestClient) -> dict[str, str]: - return get_superuser_token_headers(client) - - -@pytest.fixture(scope="module") -def normal_user_token_headers(client: TestClient, db: Session) -> dict[str, str]: - return authentication_token_from_email( - client=client, email=settings.EMAIL_TEST_USER, db=db +from app.tests.utils.utils import get_user_from_api_key +from app.models.collection import CollectionStatus + +client = TestClient(app) + +original_api_key = "ApiKey No3x47A5qoIGhm0kVKjQ77dhCqEdWRIQZlEPzzzh7i8" + + +def create_collection( + db, + user, + status: CollectionStatus = CollectionStatus.processing, + with_llm: bool = False, +): + now = datetime.now(timezone.utc) + collection = Collection( + id=uuid4(), + owner_id=user.user_id, + organization_id=user.organization_id, + project_id=user.project_id, + status=status, + updated_at=now, + ) + if with_llm: + collection.llm_service_id = f"asst_{uuid4()}" + collection.llm_service_name = "gpt-4o" + + db.add(collection) + db.commit() + db.refresh(collection) + return collection + + +def test_collection_info_processing(db: Session, client: TestClient): + headers = {"X-API-KEY": original_api_key} + user = get_user_from_api_key(db, headers) + collection = create_collection(db, user, status=CollectionStatus.processing) + + response = client.post( + f"{settings.API_V1_STR}/collections/info/{collection.id}", + headers=headers, + ) + + assert response.status_code == 200 + data = response.json()["data"] + + assert data["id"] == str(collection.id) + assert data["status"] == CollectionStatus.processing.value + assert data["llm_service_id"] is None + assert data["llm_service_name"] is None + + +def test_collection_info_successful(db: Session, client: TestClient): + headers = {"X-API-KEY": original_api_key} + user = get_user_from_api_key(db, headers) + collection = create_collection( + db, user, status=CollectionStatus.successful, with_llm=True + ) + + response = client.post( + f"{settings.API_V1_STR}/collections/info/{collection.id}", + headers=headers, + ) + + assert response.status_code == 200 + data = response.json()["data"] + + assert data["id"] == str(collection.id) + assert data["status"] == CollectionStatus.successful.value + assert data["llm_service_id"] == collection.llm_service_id + assert data["llm_service_name"] == "gpt-4o" + + +def test_collection_info_failed(db: Session, client: TestClient): + headers = {"X-API-KEY": original_api_key} + user = get_user_from_api_key(db, headers) + collection = create_collection(db, user, status=CollectionStatus.failed) + + response = client.post( + f"{settings.API_V1_STR}/collections/info/{collection.id}", + headers=headers, ) + assert response.status_code == 200 + data = response.json()["data"] -@pytest.fixture(scope="session", autouse=True) -def load_seed_data(db): - seed_database(db) - yield + assert data["id"] == str(collection.id) + assert data["status"] == CollectionStatus.failed.value + assert data["llm_service_id"] is None + assert data["llm_service_name"] is None From aa94be1a8b5aece5957c0f37b590ba025b500fca Mon Sep 17 00:00:00 2001 From: nishika26 Date: Thu, 10 Jul 2025 15:01:04 +0530 Subject: [PATCH 2/2] confest and collections test --- .../collections/test_collection_info.py | 6 +- backend/app/tests/conftest.py | 122 ++++++------------ 2 files changed, 46 insertions(+), 82 deletions(-) diff --git a/backend/app/tests/api/routes/collections/test_collection_info.py b/backend/app/tests/api/routes/collections/test_collection_info.py index 762a6654..f6ada599 100644 --- a/backend/app/tests/api/routes/collections/test_collection_info.py +++ b/backend/app/tests/api/routes/collections/test_collection_info.py @@ -38,7 +38,7 @@ def create_collection( return collection -def test_collection_info_processing(db: Session): +def test_collection_info_processing(db: Session, client: TestClient): headers = {"X-API-KEY": original_api_key} user = get_user_from_api_key(db, headers) collection = create_collection(db, user, status=CollectionStatus.processing) @@ -57,7 +57,7 @@ def test_collection_info_processing(db: Session): assert data["llm_service_name"] is None -def test_collection_info_successful(db: Session): +def test_collection_info_successful(db: Session, client: TestClient): headers = {"X-API-KEY": original_api_key} user = get_user_from_api_key(db, headers) collection = create_collection( @@ -78,7 +78,7 @@ def test_collection_info_successful(db: Session): assert data["llm_service_name"] == "gpt-4o" -def test_collection_info_failed(db: Session): +def test_collection_info_failed(db: Session, client: TestClient): headers = {"X-API-KEY": original_api_key} user = get_user_from_api_key(db, headers) collection = create_collection(db, user, status=CollectionStatus.failed) diff --git a/backend/app/tests/conftest.py b/backend/app/tests/conftest.py index f6ada599..5d68c3f1 100644 --- a/backend/app/tests/conftest.py +++ b/backend/app/tests/conftest.py @@ -1,97 +1,61 @@ -from uuid import uuid4 -from datetime import datetime, timezone +from collections.abc import Generator + +import pytest from fastapi.testclient import TestClient from sqlmodel import Session +from sqlalchemy import event + from app.core.config import settings -from app.models import Collection +from app.core.db import engine +from app.api.deps import get_db from app.main import app -from app.tests.utils.utils import get_user_from_api_key -from app.models.collection import CollectionStatus +from app.tests.utils.user import authentication_token_from_email +from app.tests.utils.utils import get_superuser_token_headers +from app.seed_data.seed_data import seed_database -client = TestClient(app) -original_api_key = "ApiKey No3x47A5qoIGhm0kVKjQ77dhCqEdWRIQZlEPzzzh7i8" +@pytest.fixture(scope="function") +def db() -> Generator[Session, None, None]: + connection = engine.connect() + transaction = connection.begin() + session = Session(bind=connection) + nested = session.begin_nested() -def create_collection( - db, - user, - status: CollectionStatus = CollectionStatus.processing, - with_llm: bool = False, -): - now = datetime.now(timezone.utc) - collection = Collection( - id=uuid4(), - owner_id=user.user_id, - organization_id=user.organization_id, - project_id=user.project_id, - status=status, - updated_at=now, - ) - if with_llm: - collection.llm_service_id = f"asst_{uuid4()}" - collection.llm_service_name = "gpt-4o" + @event.listens_for(session, "after_transaction_end") + def restart_savepoint(sess, trans): + if trans.nested and not trans._parent.nested: + sess.begin_nested() - db.add(collection) - db.commit() - db.refresh(collection) - return collection + try: + yield session + finally: + session.close() + transaction.rollback() + connection.close() -def test_collection_info_processing(db: Session, client: TestClient): - headers = {"X-API-KEY": original_api_key} - user = get_user_from_api_key(db, headers) - collection = create_collection(db, user, status=CollectionStatus.processing) +@pytest.fixture(scope="session", autouse=True) +def seed_baseline(): + with Session(engine) as session: + seed_database(session) # deterministic baseline + yield - response = client.post( - f"{settings.API_V1_STR}/collections/info/{collection.id}", - headers=headers, - ) - assert response.status_code == 200 - data = response.json()["data"] +@pytest.fixture(scope="function") +def client(db: Session): + app.dependency_overrides[get_db] = lambda: db + with TestClient(app) as c: + yield c - assert data["id"] == str(collection.id) - assert data["status"] == CollectionStatus.processing.value - assert data["llm_service_id"] is None - assert data["llm_service_name"] is None +@pytest.fixture(scope="function") +def superuser_token_headers(client: TestClient) -> dict[str, str]: + return get_superuser_token_headers(client) -def test_collection_info_successful(db: Session, client: TestClient): - headers = {"X-API-KEY": original_api_key} - user = get_user_from_api_key(db, headers) - collection = create_collection( - db, user, status=CollectionStatus.successful, with_llm=True - ) - response = client.post( - f"{settings.API_V1_STR}/collections/info/{collection.id}", - headers=headers, +@pytest.fixture(scope="function") +def normal_user_token_headers(client: TestClient, db: Session) -> dict[str, str]: + return authentication_token_from_email( + client=client, email=settings.EMAIL_TEST_USER, db=db ) - - assert response.status_code == 200 - data = response.json()["data"] - - assert data["id"] == str(collection.id) - assert data["status"] == CollectionStatus.successful.value - assert data["llm_service_id"] == collection.llm_service_id - assert data["llm_service_name"] == "gpt-4o" - - -def test_collection_info_failed(db: Session, client: TestClient): - headers = {"X-API-KEY": original_api_key} - user = get_user_from_api_key(db, headers) - collection = create_collection(db, user, status=CollectionStatus.failed) - - response = client.post( - f"{settings.API_V1_STR}/collections/info/{collection.id}", - headers=headers, - ) - - assert response.status_code == 200 - data = response.json()["data"] - - assert data["id"] == str(collection.id) - assert data["status"] == CollectionStatus.failed.value - assert data["llm_service_id"] is None - assert data["llm_service_name"] is None