Optimize Python DB tests (#3755)

* Drop db parametrization for test_returns_40x_if_not_authenticated tests

* Replace BaseModel.metadata.drop_all with _truncate_postgres_db on postgres

* Run BaseModel.metadata.create_all once on Postgres

* Drop fixtures tests

* Comment on db fixtures

* Drop more redundant test_db parametrization

* Drop skipped test

Author: r4victor

src/dstack/_internal/server/testing/conf.py

@@ -8,13 +8,17 @@
 from dstack._internal.server.db import Database, override_db
 from dstack._internal.server.models import BaseModel
 
+# Remember initialized URLs to create metadata once per session.
+_initialized_postgres_db_urls = set()
+
 
 @pytest.fixture(scope="session")
 def postgres_container():
     with PostgresContainer("postgres:16-alpine", driver="asyncpg") as postgres:
         yield postgres.get_connection_url()
 
 
+# test_db is function-scoped since making it session-scoped did not bring much benefit.
 @pytest_asyncio.fixture
 async def test_db(request):
     db_type = getattr(request, "param", "sqlite")
@@ -37,12 +41,17 @@ async def test_db(request):
         raise ValueError(f"Unknown db_type {db_type}")
     db = Database(db_url, engine=engine)
     override_db(db)
-    async with db.engine.begin() as conn:
-        await conn.run_sync(BaseModel.metadata.drop_all)
-        await conn.run_sync(BaseModel.metadata.create_all)
+    if db_type == "sqlite":
+        async with db.engine.begin() as conn:
+            await conn.run_sync(BaseModel.metadata.create_all)
+            # Relying on function-scoped engine for a clean DB
+    else:
+        if db_url not in _initialized_postgres_db_urls:
+            async with db.engine.begin() as conn:
+                await conn.run_sync(BaseModel.metadata.create_all)
+            _initialized_postgres_db_urls.add(db_url)
+        await _truncate_postgres_db(db)
     yield db
-    async with db.engine.begin() as conn:
-        await conn.run_sync(BaseModel.metadata.drop_all)
     await db.engine.dispose()
 
 
@@ -51,3 +60,15 @@ async def session(test_db):
     db = test_db
     async with db.get_session() as session:
         yield session
+
+
+async def _truncate_postgres_db(db: Database):
+    preparer = db.engine.sync_engine.dialect.identifier_preparer
+    table_names = ", ".join(
+        preparer.format_table(table) for table in BaseModel.metadata.sorted_tables
+    )
+    if not table_names:
+        return
+    truncate_statement = f"TRUNCATE {table_names} RESTART IDENTITY CASCADE"
+    async with db.engine.begin() as conn:
+        await conn.exec_driver_sql(truncate_statement)

src/tests/_internal/server/conftest.py

@@ -10,7 +10,11 @@
 from dstack._internal.server.services import logs as logs_services
 from dstack._internal.server.services.docker import ImageConfig, ImageConfigObject
 from dstack._internal.server.services.logs.filelog import FileLogStorage
-from dstack._internal.server.testing.conf import postgres_container, session, test_db  # noqa: F401
+from dstack._internal.server.testing.conf import (  # noqa: F401
+    postgres_container,
+    session,
+    test_db,
+)
 
 
 @pytest.fixture

src/tests/_internal/server/routers/test_fleets.py

@@ -59,10 +59,7 @@
 
 class TestListFleets:
     @pytest.mark.asyncio
-    @pytest.mark.parametrize("test_db", ["sqlite", "postgres"], indirect=True)
-    async def test_returns_40x_if_not_authenticated(
-        self, test_db, session: AsyncSession, client: AsyncClient
-    ):
+    async def test_returns_40x_if_not_authenticated(self, client: AsyncClient):
         response = await client.post("/api/fleets/list")
         assert response.status_code in [401, 403]
 
@@ -365,10 +362,7 @@ async def test_returns_fleet_once_if_imported_twice(
 
 class TestListProjectFleets:
     @pytest.mark.asyncio
-    @pytest.mark.parametrize("test_db", ["sqlite", "postgres"], indirect=True)
-    async def test_returns_40x_if_not_authenticated(
-        self, test_db, session: AsyncSession, client: AsyncClient
-    ):
+    async def test_returns_40x_if_not_authenticated(self, client: AsyncClient):
         response = await client.post("/api/project/main/fleets/list")
         assert response.status_code in [401, 403]
 
@@ -555,10 +549,7 @@ async def test_returns_fleet_once_if_imported_twice(
 
 class TestGetFleet:
     @pytest.mark.asyncio
-    @pytest.mark.parametrize("test_db", ["sqlite", "postgres"], indirect=True)
-    async def test_returns_40x_if_not_authenticated(
-        self, test_db, session: AsyncSession, client: AsyncClient
-    ):
+    async def test_returns_40x_if_not_authenticated(self, client: AsyncClient):
         response = await client.post("/api/project/main/fleets/get")
         assert response.status_code in [401, 403]
 
@@ -913,10 +904,7 @@ async def test_patches_profile_fleets_for_old_clients(
 
 class TestApplyFleetPlan:
     @pytest.mark.asyncio
-    @pytest.mark.parametrize("test_db", ["sqlite", "postgres"], indirect=True)
-    async def test_returns_40x_if_not_authenticated(
-        self, test_db, session: AsyncSession, client: AsyncClient
-    ):
+    async def test_returns_40x_if_not_authenticated(self, client: AsyncClient):
         response = await client.post("/api/project/main/fleets/apply")
         assert response.status_code in [401, 403]
 
@@ -1547,10 +1535,7 @@ async def test_importer_member_cannot_apply_plan_on_imported_fleet(
 
 class TestDeleteFleets:
     @pytest.mark.asyncio
-    @pytest.mark.parametrize("test_db", ["sqlite", "postgres"], indirect=True)
-    async def test_returns_40x_if_not_authenticated(
-        self, test_db, session: AsyncSession, client: AsyncClient
-    ):
+    async def test_returns_40x_if_not_authenticated(self, client: AsyncClient):
         response = await client.post("/api/project/main/fleets/delete")
         assert response.status_code in [401, 403]
 
@@ -1727,10 +1712,7 @@ async def test_importer_member_cannot_delete_imported_fleet(
 
 class TestDeleteFleetInstances:
     @pytest.mark.asyncio
-    @pytest.mark.parametrize("test_db", ["sqlite", "postgres"], indirect=True)
-    async def test_returns_40x_if_not_authenticated(
-        self, test_db, session: AsyncSession, client: AsyncClient
-    ):
+    async def test_returns_40x_if_not_authenticated(self, client: AsyncClient):
         response = await client.post("/api/project/main/fleets/delete_instances")
         assert response.status_code in [401, 403]
 
@@ -1973,10 +1955,7 @@ async def test_importer_member_cannot_delete_imported_fleet_instances(
 
 class TestGetPlan:
     @pytest.mark.asyncio
-    @pytest.mark.parametrize("test_db", ["sqlite", "postgres"], indirect=True)
-    async def test_returns_40x_if_not_authenticated(
-        self, test_db, session: AsyncSession, client: AsyncClient
-    ):
+    async def test_returns_40x_if_not_authenticated(self, client: AsyncClient):
         response = await client.post("/api/project/main/fleets/get_plan")
         assert response.status_code in [401, 403]
 

src/tests/_internal/server/routers/test_gateways.py

@@ -22,10 +22,7 @@
 
 class TestListAndGetGateways:
     @pytest.mark.asyncio
-    @pytest.mark.parametrize("test_db", ["sqlite", "postgres"], indirect=True)
-    async def test_returns_40x_if_not_authenticated(
-        self, test_db, session: AsyncSession, client: AsyncClient
-    ):
+    async def test_returns_40x_if_not_authenticated(self, client: AsyncClient):
         response = await client.post("/api/project/main/gateways/list")
         assert response.status_code in [401, 403]
 

src/tests/_internal/server/routers/test_projects.py

@@ -29,8 +29,7 @@
 
 class TestListProjects:
     @pytest.mark.asyncio
-    @pytest.mark.parametrize("test_db", ["sqlite", "postgres"], indirect=True)
-    async def test_returns_40x_if_not_authenticated(self, test_db, client: AsyncClient):
+    async def test_returns_40x_if_not_authenticated(self, client: AsyncClient):
         response = await client.post("/api/projects/list")
         assert response.status_code in [401, 403]
 
@@ -385,10 +384,7 @@ async def test_returns_total_count(self, test_db, session: AsyncSession, client:
 
 class TestListOnlyNoFleets:
     @pytest.mark.asyncio
-    @pytest.mark.parametrize("test_db", ["sqlite", "postgres"], indirect=True)
-    async def test_list_only_no_fleets_returns_40x_if_not_authenticated(
-        self, test_db, client: AsyncClient
-    ):
+    async def test_list_only_no_fleets_returns_40x_if_not_authenticated(self, client: AsyncClient):
         response = await client.post("/api/projects/list_only_no_fleets")
         assert response.status_code in [401, 403]
 
@@ -926,8 +922,7 @@ async def test_only_no_fleets_admin_requires_membership(
 
 class TestCreateProject:
     @pytest.mark.asyncio
-    @pytest.mark.parametrize("test_db", ["sqlite", "postgres"], indirect=True)
-    async def test_returns_40x_if_not_authenticated(self, test_db, client: AsyncClient):
+    async def test_returns_40x_if_not_authenticated(self, client: AsyncClient):
         response = await client.post("/api/projects/create")
         assert response.status_code in [401, 403]
 
@@ -1162,8 +1157,7 @@ async def test_creates_private_project_explicitly(
 
 class TestDeleteProject:
     @pytest.mark.asyncio
-    @pytest.mark.parametrize("test_db", ["sqlite", "postgres"], indirect=True)
-    async def test_returns_40x_if_not_authenticated(self, test_db, client: AsyncClient):
+    async def test_returns_40x_if_not_authenticated(self, client: AsyncClient):
         response = await client.post("/api/projects/delete")
         assert response.status_code in [401, 403]
 
@@ -1375,8 +1369,7 @@ async def test_errors_if_project_has_active_volumes(
 
 class TestGetProject:
     @pytest.mark.asyncio
-    @pytest.mark.parametrize("test_db", ["sqlite", "postgres"], indirect=True)
-    async def test_returns_40x_if_not_authenticated(self, test_db, client: AsyncClient):
+    async def test_returns_40x_if_not_authenticated(self, client: AsyncClient):
         response = await client.post("/api/projects/test_project/get")
         assert response.status_code in [401, 403]
 
@@ -1607,8 +1600,7 @@ async def test_member_can_access_both_public_and_private_projects(
 
 class TestSetProjectMembers:
     @pytest.mark.asyncio
-    @pytest.mark.parametrize("test_db", ["sqlite", "postgres"], indirect=True)
-    async def test_returns_40x_if_not_authenticated(self, test_db, client: AsyncClient):
+    async def test_returns_40x_if_not_authenticated(self, client: AsyncClient):
         response = await client.post("/api/projects/test_project/get")
         assert response.status_code in [401, 403]
 
@@ -1971,8 +1963,7 @@ async def test_cannot_add_same_user_twice(
 
 class TestUpdateProjectVisibility:
     @pytest.mark.asyncio
-    @pytest.mark.parametrize("test_db", ["sqlite", "postgres"], indirect=True)
-    async def test_returns_40x_if_not_authenticated(self, test_db, client: AsyncClient):
+    async def test_returns_40x_if_not_authenticated(self, client: AsyncClient):
         response = await client.post("/api/projects/test/update")
         assert response.status_code in [401, 403]
 

src/tests/_internal/server/routers/test_public_keys.py

@@ -18,13 +18,13 @@
 
 
 @pytest.mark.asyncio
-@pytest.mark.parametrize("test_db", ["sqlite", "postgres"], indirect=True)
-@pytest.mark.usefixtures("test_db")
 class TestListUserPublicKeys:
     async def test_returns_40x_if_not_authenticated(self, client: AsyncClient):
         response = await client.post("/api/users/public_keys/list")
         assert response.status_code in [401, 403]
 
+    @pytest.mark.parametrize("test_db", ["sqlite", "postgres"], indirect=True)
+    @pytest.mark.usefixtures("test_db")
     async def test_lists_own_public_keys(self, session: AsyncSession, client: AsyncClient):
         user = await create_user(session=session)
         key = await create_user_public_key(
@@ -50,6 +50,8 @@ async def test_lists_own_public_keys(self, session: AsyncSession, client: AsyncC
             }
         ]
 
+    @pytest.mark.parametrize("test_db", ["sqlite", "postgres"], indirect=True)
+    @pytest.mark.usefixtures("test_db")
     async def test_does_not_list_other_users_keys(
         self, session: AsyncSession, client: AsyncClient
     ):
@@ -63,6 +65,8 @@ async def test_does_not_list_other_users_keys(
         assert response.status_code == 200
         assert response.json() == []
 
+    @pytest.mark.parametrize("test_db", ["sqlite", "postgres"], indirect=True)
+    @pytest.mark.usefixtures("test_db")
     async def test_returns_keys_in_reverse_chronological_order(
         self, session: AsyncSession, client: AsyncClient
     ):
@@ -93,8 +97,6 @@ async def test_returns_keys_in_reverse_chronological_order(
 
 
 @pytest.mark.asyncio
-@pytest.mark.parametrize("test_db", ["sqlite", "postgres"], indirect=True)
-@pytest.mark.usefixtures("test_db")
 class TestAddUserPublicKey:
     PUBLIC_KEY_NO_COMMENT = "ssh-ed25519 AAAAC3NzaC1lZDI1NTE5AAAA"
     PUBLIC_KEY = f"{PUBLIC_KEY_NO_COMMENT} test@example.com"
@@ -114,6 +116,8 @@ async def test_returns_40x_if_not_authenticated(self, client: AsyncClient):
         )
         assert response.status_code in [401, 403]
 
+    @pytest.mark.parametrize("test_db", ["sqlite", "postgres"], indirect=True)
+    @pytest.mark.usefixtures("test_db")
     @freeze_time(datetime(2023, 1, 2, 3, 4, tzinfo=timezone.utc))
     async def test_adds_valid_public_key(
         self,
@@ -137,6 +141,8 @@ async def test_adds_valid_public_key(
         }
         validate_openssh_public_key_mock.assert_awaited_once_with(self.PUBLIC_KEY)
 
+    @pytest.mark.parametrize("test_db", ["sqlite", "postgres"], indirect=True)
+    @pytest.mark.usefixtures("test_db")
     @pytest.mark.usefixtures("validate_openssh_public_key_mock")
     async def test_adds_key_with_custom_name(self, session: AsyncSession, client: AsyncClient):
         user = await create_user(session=session)
@@ -148,6 +154,8 @@ async def test_adds_key_with_custom_name(self, session: AsyncSession, client: As
         assert response.status_code == 200
         assert response.json()["name"] == "my-laptop"
 
+    @pytest.mark.parametrize("test_db", ["sqlite", "postgres"], indirect=True)
+    @pytest.mark.usefixtures("test_db")
     @pytest.mark.usefixtures("validate_openssh_public_key_mock")
     async def test_uses_md5_as_name_when_no_comment_and_no_name(
         self, session: AsyncSession, client: AsyncClient
@@ -161,6 +169,8 @@ async def test_uses_md5_as_name_when_no_comment_and_no_name(
         assert response.status_code == 200
         assert response.json()["name"] == "744e414c6ac55e3f15c1dd48229cbe74"
 
+    @pytest.mark.parametrize("test_db", ["sqlite", "postgres"], indirect=True)
+    @pytest.mark.usefixtures("test_db")
     @pytest.mark.parametrize(
         "key",
         [
@@ -180,6 +190,8 @@ async def test_returns_400_for_invalid_key(
         assert response.status_code == 400
         assert "Invalid public key" in response.json()["detail"][0]["msg"]
 
+    @pytest.mark.parametrize("test_db", ["sqlite", "postgres"], indirect=True)
+    @pytest.mark.usefixtures("test_db")
     async def test_returns_400_for_unsupported_key(
         self, session: AsyncSession, client: AsyncClient
     ):
@@ -192,6 +204,8 @@ async def test_returns_400_for_unsupported_key(
         assert response.status_code == 400
         assert response.json()["detail"][0]["msg"] == "Unsupported key type: ssh-dss"
 
+    @pytest.mark.parametrize("test_db", ["sqlite", "postgres"], indirect=True)
+    @pytest.mark.usefixtures("test_db")
     @pytest.mark.usefixtures("validate_openssh_public_key_mock")
     async def test_returns_400_resource_exists_for_duplicate_key(
         self, session: AsyncSession, client: AsyncClient
@@ -214,8 +228,6 @@ async def test_returns_400_resource_exists_for_duplicate_key(
 
 
 @pytest.mark.asyncio
-@pytest.mark.parametrize("test_db", ["sqlite", "postgres"], indirect=True)
-@pytest.mark.usefixtures("test_db")
 class TestDeleteUserPublicKeys:
     async def test_returns_40x_if_not_authenticated(self, client: AsyncClient):
         response = await client.post(
@@ -224,6 +236,8 @@ async def test_returns_40x_if_not_authenticated(self, client: AsyncClient):
         )
         assert response.status_code in [401, 403]
 
+    @pytest.mark.parametrize("test_db", ["sqlite", "postgres"], indirect=True)
+    @pytest.mark.usefixtures("test_db")
     async def test_deletes_public_key(self, session: AsyncSession, client: AsyncClient):
         user = await create_user(session=session)
         key = await create_user_public_key(session=session, user=user)
@@ -241,6 +255,8 @@ async def test_deletes_public_key(self, session: AsyncSession, client: AsyncClie
         )
         assert res.scalars().all() == [other_key]
 
+    @pytest.mark.parametrize("test_db", ["sqlite", "postgres"], indirect=True)
+    @pytest.mark.usefixtures("test_db")
     async def test_silently_ignores_nonexistent_ids(
         self, session: AsyncSession, client: AsyncClient
     ):
@@ -252,6 +268,8 @@ async def test_silently_ignores_nonexistent_ids(
         )
         assert response.status_code == 200
 
+    @pytest.mark.parametrize("test_db", ["sqlite", "postgres"], indirect=True)
+    @pytest.mark.usefixtures("test_db")
     async def test_does_not_delete_other_users_keys(
         self, session: AsyncSession, client: AsyncClient
     ):

src/tests/_internal/server/routers/test_runs.py

@@ -642,10 +642,7 @@ def get_service_run_spec(
 
 class TestListRuns:
     @pytest.mark.asyncio
-    @pytest.mark.parametrize("test_db", ["sqlite", "postgres"], indirect=True)
-    async def test_returns_40x_if_not_authenticated(
-        self, test_db, session: AsyncSession, client: AsyncClient
-    ):
+    async def test_returns_40x_if_not_authenticated(self, client: AsyncClient):
         response = await client.post("/api/runs/list")
         assert response.status_code in [401, 403]
 

src/tests/_internal/server/routers/test_server.py

@@ -8,7 +8,7 @@
 
 class TestGetInfo:
     @pytest.mark.asyncio
-    async def test_returns_40x_if_not_authenticated(self, test_db, client: AsyncClient):
+    async def test_returns_server_info(self, test_db, client: AsyncClient):
         with patch.object(settings, "DSTACK_VERSION", "0.18.10"):
             response = await client.post("/api/server/get_info")
         assert response.status_code == 200