Fix integration and unit tests (#912)

* WIP

* WIP

* wip

* Conftest revert

* Add default database for testing

* TLK-2723 Resolve current integration/unit test issues

* TLK-2723 Resolve current integration/unit test issues

* TLK-2723 Resolve current integration/unit test issues

---------

Co-authored-by: EugeneP <[email protected]>

Makefile

@@ -56,7 +56,7 @@ run-community-tests-debug:
 
 .PHONY: run-integration-tests
 run-integration-tests:
-	docker compose run --rm --build backend poetry run pytest -c src/backend/pytest_integration.ini src/backend/tests/integration/$(file)
+	poetry run pytest -c src/backend/pytest_integration.ini src/backend/tests/integration/$(file)
 
 .PHONY: test-db
 test-db:

src/backend/alembic/versions/2024_08_01_117f0d9b1d3d_seed_deployments_data.py

@@ -10,9 +10,9 @@
 
 from alembic import op
 
-from backend.database_models.seeders.deployments_models_seed import (
-    delete_default_models,
-    deployments_models_seed,
+from backend.database_models.seeders.organization_seed import (
+    delete_default_organization,
+    seed_default_organization,
 )
 
 # revision identifiers, used by Alembic.
@@ -23,8 +23,8 @@
 
 
 def upgrade() -> None:
-    deployments_models_seed(op)
+    seed_default_organization(op)
 
 
 def downgrade() -> None:
-    delete_default_models(op)
+    delete_default_organization(op)

src/backend/chat/custom/utils.py

@@ -1,7 +1,6 @@
 from typing import Any
 
 from backend.database_models.database import get_session
-from backend.exceptions import DeploymentNotFoundError
 from backend.model_deployments.base import BaseDeployment
 from backend.schemas.context import Context
 from backend.services import deployment as deployment_service
@@ -20,7 +19,7 @@ def get_deployment(name: str, ctx: Context, **kwargs: Any) -> BaseDeployment:
     try:
         session = next(get_session())
         deployment = deployment_service.get_deployment_by_name(session, name, **kwargs)
-    except DeploymentNotFoundError:
+    except Exception:
         deployment = deployment_service.get_default_deployment(**kwargs)
 
     return deployment

src/backend/database_models/seeders/organization_seed.py

@@ -0,0 +1,42 @@
+from sqlalchemy import text
+from sqlalchemy.orm import Session
+
+from backend.database_models import Organization
+
+
+def seed_default_organization(op):
+    """
+    Seed default organization.
+    """
+    # Previously we would seed the default deployments and models here. We've changed this
+    # behaviour during a refactor of the deployments module so that deployments and models
+    # are inserted when they're first used. This solves an issue where seed data would
+    # sometimes be inserted with invalid config data.
+
+    _ = Session(op.get_bind())
+
+    # Seed default organization
+    sql_command = text(
+        """
+        INSERT INTO organizations (
+            id, name, created_at, updated_at
+        )
+        VALUES (
+            :id, :name, now(), now()
+        )
+        ON CONFLICT (id) DO NOTHING;
+    """
+    ).bindparams(
+        id="default",
+        name="Default Organization",
+    )
+    op.execute(sql_command)
+
+
+def delete_default_organization(op):
+    """
+    Delete default organization.
+    """
+    session = Session(op.get_bind())
+    session.query(Organization).filter_by(id="default").delete()
+    session.commit()

src/backend/pytest_integration.ini

@@ -1,5 +1,6 @@
 [pytest]
 env =
-    DATABASE_URL=postgresql://postgres:postgres@db:5432/postgres
+    DATABASE_URL=postgresql://postgres:postgres@localhost:5433
 filterwarnings =
     ignore::UserWarning:pydantic.*
+    ignore::DeprecationWarning

src/backend/services/deployment.py

@@ -91,7 +91,7 @@ def get_deployment_definition_by_name(session: DBSessionDep, deployment_name: st
         raise DeploymentNotFoundError(deployment_id=deployment_name)
 
     if definition.name not in [d.name for d in deployment_crud.get_deployments(session)]:
-        create_db_deployment(session, definition)
+        definition = create_db_deployment(session, definition)
 
     return definition
 

src/backend/services/request_validators.py

@@ -11,6 +11,7 @@
 from backend.crud import model as model_crud
 from backend.crud import organization as organization_crud
 from backend.database_models.database import DBSessionDep
+from backend.exceptions import DeploymentNotFoundError
 from backend.model_deployments.utils import class_name_validator
 from backend.services import deployment as deployment_service
 from backend.services.agent import validate_agent_exists
@@ -217,7 +218,12 @@ async def validate_env_vars(session: DBSessionDep, request: Request):
     invalid_keys = []
 
     deployment_id = unquote_plus(request.path_params.get("deployment_id"))
-    deployment = deployment_service.get_deployment(session, deployment_id)
+    try:
+        deployment = deployment_service.get_deployment(session, deployment_id)
+    except DeploymentNotFoundError:
+        raise HTTPException(
+            status_code=404, detail=f"Deployment {deployment_id} not found."
+        )
 
     for key in env_vars:
         if key not in deployment.env_vars():

src/backend/tests/integration/conftest.py

@@ -18,7 +18,7 @@
 from backend.schemas.user import User
 from backend.tests.unit.factories import get_factory
 
-DATABASE_URL = os.environ["DATABASE_URL"]
+DATABASE_URL = os.environ.get("DATABASE_URL", "postgresql://postgres:postgres@localhost:5433")
 
 
 @pytest.fixture
@@ -162,7 +162,6 @@ def deployment(session: Session) -> Deployment:
         deployment_class_name="CohereDeployment"
     )
 
-
 @pytest.fixture
 def model(session: Session) -> Model:
     return get_factory("Model", session).create()

src/backend/tests/integration/routers/test_agent.py

@@ -283,7 +283,8 @@ def test_create_agent_deployment_not_in_db(
         "deployment": CohereDeployment.name(),
     }
     cohere_deployment = deployment_crud.get_deployment_by_name(session, CohereDeployment.name())
-    deployment_crud.delete_deployment(session, cohere_deployment.id)
+    if cohere_deployment:
+        deployment_crud.delete_deployment(session, cohere_deployment.id)
     response = session_client.post(
         "/v1/agents", json=request_json, headers={"User-Id": user.id}
     )

src/backend/tests/integration/routers/test_conversation.py

@@ -11,6 +11,10 @@
 from backend.tests.unit.factories import get_factory
 
 
+@pytest.mark.skipif(
+    os.environ.get("COHERE_API_KEY") is None,
+    reason="Cohere API key not set, skipping test",
+)
 def test_search_conversations(
     session_client: TestClient,
     session: Session,
@@ -64,7 +68,10 @@ def test_search_conversations_with_reranking(
     assert len(results) == 1
     assert results[0]["id"] == conversation2.id
 
-
+@pytest.mark.skipif(
+    os.environ.get("COHERE_API_KEY") is None,
+    reason="Cohere API key not set, skipping test",
+)
 def test_search_conversations_no_conversations(
     session_client: TestClient,
     session: Session,

src/backend/tests/unit/configuration.yaml

@@ -1,5 +1,5 @@
 deployments:
-  default_deployment:
+  default_deployment: cohere_platform
   enabled_deployments:
   sagemaker:
     access_key: "sagemaker_access_key"

src/backend/tests/unit/conftest.py

@@ -15,12 +15,13 @@
 
 from backend.database_models import get_session
 from backend.database_models.base import CustomFilterQuery
+from backend.database_models.deployment import Deployment
 from backend.main import app, create_app
 from backend.schemas.organization import Organization
 from backend.schemas.user import User
 from backend.tests.unit.factories import get_factory
 
-DATABASE_URL = os.environ["DATABASE_URL"]
+DATABASE_URL = os.environ.get("DATABASE_URL", "postgresql://postgres:postgres@localhost:5433")
 MASTER_DB_NAME = "postgres"
 TEST_DB_PREFIX = "postgres_"
 MASTER_DATABASE_FULL_URL = f"{DATABASE_URL}/{MASTER_DB_NAME}"
@@ -58,7 +59,7 @@ def client():
     yield TestClient(app)
 
 
-@pytest.fixture(scope="session")
+@pytest.fixture
 def engine(worker_id: str) -> Generator[Any, None, None]:
     """
     Yields a SQLAlchemy engine which is disposed of after the test session
@@ -81,6 +82,32 @@ def engine(worker_id: str) -> Generator[Any, None, None]:
     drop_test_database_if_exists(test_db_name)
 
 
+@pytest.fixture(scope="session")
+def engine_chat(worker_id: str) -> Generator[Any, None, None]:
+    """
+    Yields a SQLAlchemy engine which is disposed of after the test session
+    """
+    test_db_name = f"{TEST_DB_PREFIX}{worker_id}"
+    if worker_id == "master":
+        test_db_name = f"{TEST_DB_PREFIX}{worker_id}_chat"
+
+    test_db_url = f"{DATABASE_URL}/{test_db_name}"
+
+    drop_test_database_if_exists(test_db_name)
+    create_test_database(test_db_name)
+    engine = create_engine(test_db_url, echo=True)
+
+    with engine.begin():
+        alembic_cfg = Config("src/backend/alembic.ini")
+        alembic_cfg.set_main_option("sqlalchemy.url", test_db_url)
+        upgrade(alembic_cfg, "head")
+
+    yield engine
+
+    engine.dispose()
+    drop_test_database_if_exists(test_db_name)
+
+
 @pytest.fixture(scope="function")
 def session(engine: Any) -> Generator[Session, None, None]:
     """
@@ -122,15 +149,15 @@ def override_get_session() -> Generator[Session, Any, None]:
 
 
 @pytest.fixture(scope="session")
-def session_chat(engine: Any) -> Generator[Session, None, None]:
+def session_chat(engine_chat: Any) -> Generator[Session, None, None]:
     """
     Yields a SQLAlchemy session within a transaction
     that is rolled back after every session
 
     We need to use the fixture in the session scope because the chat
     endpoint is asynchronous and needs to be open for the entire session
     """
-    connection = engine.connect()
+    connection = engine_chat.connect()
     transaction = connection.begin()
     # Use connection within the started transaction
     session = Session(bind=connection, query_cls=CustomFilterQuery)
@@ -188,6 +215,11 @@ def user(session: Session) -> User:
 def organization(session: Session) -> Organization:
     return get_factory("Organization", session).create()
 
+@pytest.fixture
+def deployment(session: Session) -> Deployment:
+    return get_factory("Deployment", session).create(
+        deployment_class_name="CohereDeployment"
+    )
 
 @pytest.fixture
 def mock_available_model_deployments(request):

src/backend/tests/integration/routers/test_deployment.py → src/backend/tests/unit/routers/test_deployment.py

@@ -157,9 +157,9 @@ def test_set_env_vars(
 
 
 def test_set_env_vars_with_invalid_deployment_name(
-    client: TestClient
+    session_client: TestClient
 ):
-    response = client.post("/v1/deployments/unknown/update_config", json={})
+    response = session_client.post("/v1/deployments/unknown/update_config", json={})
     assert response.status_code == 404
 
 

src/backend/tests/unit/services/test_deployment.py

@@ -83,7 +83,12 @@ def test_get_deployment_definition_by_name(session, mock_available_model_deploym
 
 def test_get_deployment_definition_by_name_no_db_deployments(session, mock_available_model_deployments, clear_db_deployments) -> None:
     definition = deployment_service.get_deployment_definition_by_name(session, MockCohereDeployment.name())
-    assert definition == MockCohereDeployment.to_deployment_definition()
+    mock = MockCohereDeployment.to_deployment_definition()
+    assert definition.name == mock.name
+    assert definition.models == mock.models
+    assert definition.class_name == mock.class_name
+    assert definition.config == mock.config
+
 
 def test_get_deployment_definition_by_name_wrong_name(session, mock_available_model_deployments) -> None:
     with pytest.raises(DeploymentNotFoundError):