test(BA-6644): add session scope ids to remaining test fixtures

Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>

Author: seedspirit

tests/unit/manager/models/resource_slot/conftest.py

@@ -6,6 +6,8 @@
 
 import pytest
 
+from ai.backend.common.identifier.domain import DomainID
+from ai.backend.common.identifier.resource_group import ResourceGroupID
 from ai.backend.common.types import BinarySize, ResourceSlot
 from ai.backend.manager.data.auth.hash import PasswordHashAlgorithm
 from ai.backend.manager.models.agent import AgentRow
@@ -67,24 +69,37 @@ async def database_with_resource_slot_tables(
         yield database_connection
 
 
+@pytest.fixture
+def domain_id() -> DomainID:
+    return DomainID(uuid.uuid4())
+
+
+@pytest.fixture
+def scaling_group_id() -> ResourceGroupID:
+    return ResourceGroupID(uuid.uuid4())
+
+
 @pytest.fixture
 async def domain_name(
     database_with_resource_slot_tables: ExtendedAsyncSAEngine,
+    domain_id: DomainID,
 ) -> AsyncGenerator[str, None]:
     name = "test-domain"
     async with database_with_resource_slot_tables.begin_session() as db_sess:
-        db_sess.add(DomainRow(name=name))
+        db_sess.add(DomainRow(id=domain_id, name=name))
     yield name
 
 
 @pytest.fixture
 async def scaling_group(
     database_with_resource_slot_tables: ExtendedAsyncSAEngine,
+    scaling_group_id: ResourceGroupID,
 ) -> AsyncGenerator[str, None]:
     name = "default"
     async with database_with_resource_slot_tables.begin_session() as db_sess:
         db_sess.add(
             ScalingGroupRow(
+                id=scaling_group_id,
                 name=name,
                 description="Test scaling group",
                 is_active=True,
@@ -208,9 +223,11 @@ async def agent_id(
 @pytest.fixture
 async def kernel_id(
     database_with_resource_slot_tables: ExtendedAsyncSAEngine,
+    domain_id: DomainID,
     domain_name: str,
     project_id: uuid.UUID,
     user_uuid: uuid.UUID,
+    scaling_group_id: ResourceGroupID,
     scaling_group: str,
     agent_id: str,
 ) -> AsyncGenerator[uuid.UUID, None]:
@@ -220,8 +237,10 @@ async def kernel_id(
         db_sess.add(
             SessionRow(
                 id=sid,
+                domain_id=domain_id,
                 domain_name=domain_name,
                 group_id=project_id,
+                resource_group_id=scaling_group_id,
                 scaling_group_name=scaling_group,
                 user_uuid=user_uuid,
                 occupying_slots=ResourceSlot({"cpu": Decimal("1"), "mem": Decimal("1073741824")}),

tests/unit/manager/repositories/domain/test_domain_purgers.py

@@ -12,6 +12,7 @@
 import pytest
 import sqlalchemy as sa
 
+from ai.backend.common.identifier.resource_group import ResourceGroupID
 from ai.backend.common.types import ResourceSlot, VFolderHostPermissionMap
 from ai.backend.manager.data.auth.hash import PasswordHashAlgorithm
 from ai.backend.manager.data.kernel.types import KernelStatus
@@ -182,9 +183,11 @@ async def sample_sessions(
         """Create test sessions belonging to the domain."""
         sessions: list[SessionRow] = []
         sgroup_name = f"default-{uuid.uuid4().hex[:8]}"
+        sgroup_id = ResourceGroupID(uuid.uuid4())
         async with db_with_cleanup.begin_session() as session:
             sgroup = ScalingGroupRow(
                 name=sgroup_name,
+                id=sgroup_id,
                 description="Test scaling group",
                 is_active=True,
                 driver="static",
@@ -201,8 +204,10 @@ async def sample_sessions(
                     cluster_mode="single-node",
                     cluster_size=1,
                     domain_name=sample_domain.domain_name,
+                    domain_id=sample_domain.domain_id,
                     group_id=sample_group.id,
                     scaling_group_name=sgroup_name,
+                    resource_group_id=sgroup_id,
                     user_uuid=sample_user.uuid,
                     occupying_slots=ResourceSlot({}),
                     requested_slots=ResourceSlot({}),

tests/unit/manager/repositories/domain/test_domain_repository.py

@@ -12,7 +12,8 @@
 import sqlalchemy as sa
 
 from ai.backend.common.exception import DomainNotFound, InvalidAPIParameters
-from ai.backend.common.identifier.domain import DomainName
+from ai.backend.common.identifier.domain import DomainID, DomainName
+from ai.backend.common.identifier.resource_group import ResourceGroupID
 from ai.backend.common.types import (
     DefaultForUnspecified,
     ResourceSlot,
@@ -333,13 +334,16 @@ async def domain_with_active_kernel(
     ) -> str:
         """Create an inactive domain with an active kernel for purge testing."""
         domain_name = f"domain-with-kernel-{uuid.uuid4().hex[:8]}"
+        domain_id = DomainID(uuid.uuid4())
         user_uuid = uuid.uuid4()
         group_id = uuid.uuid4()
         session_id = uuid.uuid4()
         sgroup_name = f"default-{uuid.uuid4().hex[:8]}"
+        sgroup_id = ResourceGroupID(uuid.uuid4())
 
         async with db_with_default_resource_policies.begin_session() as session:
             domain = DomainRow(
+                id=domain_id,
                 name=domain_name,
                 description="Test domain with active kernels",
                 is_active=False,
@@ -352,6 +356,7 @@ async def domain_with_active_kernel(
             session.add(domain)
 
             sgroup = ScalingGroupRow(
+                id=sgroup_id,
                 name=sgroup_name,
                 description="Test scaling group",
                 is_active=True,
@@ -403,8 +408,10 @@ async def domain_with_active_kernel(
                 creation_id=str(uuid.uuid4()).replace("-", ""),
                 cluster_size=1,
                 domain_name=domain_name,
+                domain_id=domain_id,
                 group_id=group_id,
                 scaling_group_name=sgroup_name,
+                resource_group_id=sgroup_id,
                 user_uuid=user_uuid,
                 occupying_slots=ResourceSlot(),
                 requested_slots=ResourceSlot(),

tests/unit/manager/repositories/group/test_group_db_source.py

@@ -10,6 +10,8 @@
 import sqlalchemy as sa
 
 from ai.backend.common.identifier.deployment import DeploymentID
+from ai.backend.common.identifier.domain import DomainID
+from ai.backend.common.identifier.resource_group import ResourceGroupID
 from ai.backend.common.types import ResourceSlot, VFolderHostPermissionMap
 from ai.backend.manager.data.auth.hash import PasswordHashAlgorithm
 from ai.backend.manager.data.group.types import ProjectType
@@ -111,16 +113,22 @@ async def db_with_cleanup(
         ):
             yield database_connection
 
+    @pytest.fixture
+    def test_domain_id(self) -> DomainID:
+        return DomainID(uuid.uuid4())
+
     @pytest.fixture
     async def test_domain(
         self,
         db_with_cleanup: ExtendedAsyncSAEngine,
+        test_domain_id: DomainID,
     ) -> str:
         """Create test domain"""
         domain_name = f"test-domain-{uuid.uuid4().hex[:8]}"
 
         async with db_with_cleanup.begin_session() as session:
             domain = DomainRow(
+                id=test_domain_id,
                 name=domain_name,
                 description="Test domain",
                 is_active=True,
@@ -290,17 +298,20 @@ async def inactive_endpoint_with_session_and_routing(
         self,
         db_with_cleanup: ExtendedAsyncSAEngine,
         test_domain: str,
+        test_domain_id: DomainID,
         test_user: uuid.UUID,
         test_group: uuid.UUID,
     ) -> EndpointWithSessionFixtureData:
         """Create one inactive endpoint with a session and routing entry"""
         endpoint_id = DeploymentID(uuid.uuid4())
         session_id = uuid.uuid4()
+        sgroup_id = ResourceGroupID(uuid.uuid4())
         sgroup_name = f"default-{uuid.uuid4().hex[:8]}"
 
         async with db_with_cleanup.begin_session() as session:
             # Create scaling group
             sgroup = ScalingGroupRow(
+                id=sgroup_id,
                 name=sgroup_name,
                 description="Test scaling group",
                 is_active=True,
@@ -332,8 +343,10 @@ async def inactive_endpoint_with_session_and_routing(
                 id=session_id,
                 creation_id=f"test-session-{uuid.uuid4().hex[:8]}",
                 domain_name=test_domain,
+                domain_id=test_domain_id,
                 group_id=test_group,
                 scaling_group_name=sgroup_name,
+                resource_group_id=sgroup_id,
                 user_uuid=test_user,
                 access_key="test-access-key",
                 cluster_mode="single-node",
@@ -413,17 +426,20 @@ async def multiple_endpoints_with_sessions(
         self,
         db_with_cleanup: ExtendedAsyncSAEngine,
         test_domain: str,
+        test_domain_id: DomainID,
         test_user: uuid.UUID,
         test_group: uuid.UUID,
     ) -> MultipleEndpointsWithSessionsFixtureData:
         """Create 3 inactive endpoints, each with a session and routing entry"""
         endpoint_ids = []
         session_ids = []
+        sgroup_id = ResourceGroupID(uuid.uuid4())
         sgroup_name = f"default-{uuid.uuid4().hex[:8]}"
 
         async with db_with_cleanup.begin_session() as session:
             # Create scaling group
             sgroup = ScalingGroupRow(
+                id=sgroup_id,
                 name=sgroup_name,
                 description="Test scaling group",
                 is_active=True,
@@ -459,8 +475,10 @@ async def multiple_endpoints_with_sessions(
                     id=session_id,
                     creation_id=f"test-session-{i}-{uuid.uuid4().hex[:8]}",
                     domain_name=test_domain,
+                    domain_id=test_domain_id,
                     group_id=test_group,
                     scaling_group_name=sgroup_name,
+                    resource_group_id=sgroup_id,
                     user_uuid=test_user,
                     access_key=f"test-access-key-{i}",
                     cluster_mode="single-node",

tests/unit/manager/repositories/group/test_group_purgers.py

@@ -12,6 +12,7 @@
 import pytest
 import sqlalchemy as sa
 
+from ai.backend.common.identifier.resource_group import ResourceGroupID
 from ai.backend.common.types import ResourceSlot
 from ai.backend.manager.data.auth.hash import PasswordHashAlgorithm
 from ai.backend.manager.data.kernel.types import KernelStatus
@@ -185,12 +186,21 @@ async def sample_group(
             return group
 
     @pytest.fixture
-    async def sample_scaling_group(self, db_with_cleanup: ExtendedAsyncSAEngine) -> str:
+    def sample_scaling_group_id(self) -> ResourceGroupID:
+        return ResourceGroupID(uuid.uuid4())
+
+    @pytest.fixture
+    async def sample_scaling_group(
+        self,
+        db_with_cleanup: ExtendedAsyncSAEngine,
+        sample_scaling_group_id: ResourceGroupID,
+    ) -> str:
         """Create a test scaling group."""
         sgroup_name = f"test-sgroup-{uuid.uuid4().hex[:8]}"
         async with db_with_cleanup.begin_session() as session:
             sgroup = ScalingGroupRow(
                 name=sgroup_name,
+                id=sample_scaling_group_id,
                 driver="static",
                 driver_opts={},
                 scheduler="fifo",
@@ -207,6 +217,7 @@ async def sample_sessions(
         sample_domain: DomainFixtureData,
         sample_user: UserRow,
         sample_scaling_group: str,
+        sample_scaling_group_id: ResourceGroupID,
     ) -> list[SessionRow]:
         """Create test sessions belonging to the group."""
         sessions: list[SessionRow] = []
@@ -218,8 +229,10 @@ async def sample_sessions(
                     cluster_mode="single-node",
                     cluster_size=1,
                     domain_name=sample_domain.domain_name,
+                    domain_id=sample_domain.domain_id,
                     group_id=sample_group.id,
                     scaling_group_name=sample_scaling_group,
+                    resource_group_id=sample_scaling_group_id,
                     user_uuid=sample_user.uuid,
                     occupying_slots=ResourceSlot({}),
                     requested_slots=ResourceSlot({}),
@@ -310,6 +323,7 @@ async def sample_endpoint_sessions_with_routings(
         sample_group: GroupRow,
         sample_user: UserRow,
         sample_scaling_group: str,
+        sample_scaling_group_id: ResourceGroupID,
     ) -> tuple[list[SessionRow], list[RoutingRow]]:
         """Create test sessions connected to endpoints via routings."""
         sessions: list[SessionRow] = []
@@ -323,8 +337,10 @@ async def sample_endpoint_sessions_with_routings(
                     cluster_mode="single-node",
                     cluster_size=1,
                     domain_name=sample_domain.domain_name,
+                    domain_id=sample_domain.domain_id,
                     group_id=sample_group.id,
                     scaling_group_name=sample_scaling_group,
+                    resource_group_id=sample_scaling_group_id,
                     user_uuid=sample_user.uuid,
                     occupying_slots=ResourceSlot({}),
                     requested_slots=ResourceSlot({}),