refactor: split scaling group component fixtures

Author: seedspirit

src/ai/backend/testutils/fixtures.py

@@ -4,12 +4,10 @@
 from dataclasses import dataclass
 
 from ai.backend.common.identifier.domain import DomainID, DomainName
-from ai.backend.common.identifier.resource_group import ResourceGroupID, ResourceGroupName
 
 __all__ = (
     "DomainFactory",
     "DomainFixtureData",
-    "ScalingGroupFixtureData",
 )
 
 
@@ -19,10 +17,4 @@ class DomainFixtureData:
     domain_id: DomainID
 
 
-@dataclass(frozen=True)
-class ScalingGroupFixtureData:
-    scaling_group_name: ResourceGroupName
-    scaling_group_id: ResourceGroupID
-
-
 DomainFactory = Callable[..., Awaitable[DomainFixtureData]]

tests/component/agent_api/conftest.py

@@ -11,6 +11,7 @@
 
 from ai.backend.common.etcd import AsyncEtcd, ConfigScopes
 from ai.backend.common.events.dispatcher import EventProducer
+from ai.backend.common.identifier.resource_group import ResourceGroupName
 from ai.backend.common.plugin.hook import HookPluginContext
 from ai.backend.common.types import HostPortPair, ResourceSlot
 from ai.backend.manager.actions.validators import ActionValidators
@@ -28,7 +29,6 @@
 from ai.backend.manager.repositories.scheduler.repository import SchedulerRepository
 from ai.backend.manager.services.agent.processors import AgentProcessors
 from ai.backend.manager.services.agent.service import AgentService
-from ai.backend.testutils.fixtures import ScalingGroupFixtureData
 
 
 @pytest.fixture()
@@ -106,11 +106,11 @@ def server_module_registries(
 @pytest.fixture()
 async def agent_fixture(
     db_engine: SAEngine,
-    scaling_group_fixture: ScalingGroupFixtureData,
+    scaling_group_name: ResourceGroupName,
 ) -> AsyncIterator[str]:
     """Insert a test agent row and yield its ID.
 
-    The agent references the scaling_group_fixture via FK.
+    The agent references the scaling_group_name via FK.
     Teardown deletes the agent row (cascade deletes agent_resources).
     """
     agent_id = f"i-test-agent-{secrets.token_hex(6)}"
@@ -120,7 +120,7 @@ async def agent_fixture(
                 id=agent_id,
                 status=AgentStatus.ALIVE,
                 region="local",
-                scaling_group=scaling_group_fixture.scaling_group_name,
+                scaling_group=scaling_group_name,
                 schedulable=True,
                 available_slots=ResourceSlot({"cpu": "4", "mem": "8589934592"}),
                 occupied_slots=ResourceSlot(),

tests/component/agent_api/test_agent_api.py

@@ -26,7 +26,7 @@
     OrderDirection,
 )
 from ai.backend.common.dto.manager.query import StringFilter
-from ai.backend.testutils.fixtures import ScalingGroupFixtureData
+from ai.backend.common.identifier.resource_group import ResourceGroupName
 
 
 class TestSearchAgents:
@@ -84,10 +84,10 @@ async def test_admin_searches_agents_with_resource_group_filter(
         self,
         admin_registry: BackendAIClientRegistry,
         agent_fixture: str,
-        scaling_group_fixture: ScalingGroupFixtureData,
+        scaling_group_name: ResourceGroupName,
     ) -> None:
         """Filtering by resource_group returns agents in that scaling group."""
-        sgroup_name = scaling_group_fixture.scaling_group_name
+        sgroup_name = scaling_group_name
         result = await admin_registry.agent.search_agents(
             SearchAgentsRequest(
                 filter=AgentFilter(
@@ -140,11 +140,11 @@ async def test_compound_filters_status_and_resource_group(
         self,
         admin_registry: BackendAIClientRegistry,
         agent_fixture: str,
-        scaling_group_fixture: ScalingGroupFixtureData,
+        scaling_group_name: ResourceGroupName,
     ) -> None:
         """Compound filters (status + resource_group) return intersection of conditions."""
         # Search with both status=ALIVE and resource_group filters
-        sgroup_name = scaling_group_fixture.scaling_group_name
+        sgroup_name = scaling_group_name
         result = await admin_registry.agent.search_agents(
             SearchAgentsRequest(
                 filter=AgentFilter(

tests/component/auto_scaling_rule/conftest.py

@@ -11,6 +11,7 @@
 
 from ai.backend.common.container_registry import ContainerRegistryType
 from ai.backend.common.data.endpoint.types import EndpointLifecycle
+from ai.backend.common.identifier.resource_group import ResourceGroupName
 from ai.backend.manager.actions.validators import ActionValidators
 from ai.backend.manager.actions.validators.rbac import RBACValidators
 from ai.backend.manager.actions.validators.rbac.bulk import BulkActionRBACValidator
@@ -37,7 +38,7 @@
 )
 from ai.backend.manager.services.deployment.processors import DeploymentProcessors
 from ai.backend.manager.services.deployment.service import DeploymentService
-from ai.backend.testutils.fixtures import DomainFixtureData, ScalingGroupFixtureData
+from ai.backend.testutils.fixtures import DomainFixtureData
 
 
 @dataclass
@@ -125,7 +126,7 @@ async def model_deployment_fixture(
     db_engine: SAEngine,
     domain_fixture: DomainFixtureData,
     group_fixture: uuid.UUID,
-    scaling_group_fixture: ScalingGroupFixtureData,
+    scaling_group_name: ResourceGroupName,
     admin_user_fixture: UserFixtureData,
 ) -> AsyncIterator[uuid.UUID]:
     """Insert a minimal EndpointRow (model deployment) and yield its UUID.
@@ -176,7 +177,7 @@ async def model_deployment_fixture(
                 session_owner=str(admin_user_fixture.user_uuid),
                 domain=domain_fixture.domain_name,
                 project=str(group_fixture),
-                resource_group=scaling_group_fixture.scaling_group_name,
+                resource_group=scaling_group_name,
                 lifecycle_stage=EndpointLifecycle.CREATED,
                 url=None,
             )

tests/component/conftest.py

@@ -152,7 +152,7 @@
     postgres_container,
     redis_container,
 )
-from ai.backend.testutils.fixtures import DomainFixtureData, ScalingGroupFixtureData
+from ai.backend.testutils.fixtures import DomainFixtureData
 from ai.backend.testutils.pants import get_parallel_slot
 
 log = logging.getLogger("tests.component.conftest")
@@ -673,16 +673,15 @@ async def resource_policy_fixture(
 
 
 @pytest.fixture()
-async def scaling_group_fixture(
+async def scaling_group_name(
     db_engine: SAEngine,
     domain_fixture: DomainFixtureData,
-) -> AsyncIterator[ScalingGroupFixtureData]:
-    """Insert a scaling group and its domain association; yield its identifiers."""
+) -> AsyncIterator[ResourceGroupName]:
+    """Insert a scaling group and its domain association; yield its name."""
     sgroup_name = ResourceGroupName(f"sgroup-{secrets.token_hex(6)}")
     async with db_engine.begin() as conn:
-        result = await conn.execute(
-            sa.insert(scaling_groups)
-            .values(
+        await conn.execute(
+            sa.insert(scaling_groups).values(
                 name=sgroup_name,
                 description=f"Test scaling group {sgroup_name}",
                 is_active=True,
@@ -691,23 +690,34 @@ async def scaling_group_fixture(
                 scheduler="fifo",
                 scheduler_opts=ScalingGroupOpts(),
             )
-            .returning(scaling_groups.c.id)
         )
-        sgroup_id = ResourceGroupID(result.scalar_one())
         await conn.execute(
             sa.insert(sgroups_for_domains).values(
                 scaling_group=sgroup_name,
                 domain=domain_fixture.domain_name,
             )
         )
-    yield ScalingGroupFixtureData(scaling_group_name=sgroup_name, scaling_group_id=sgroup_id)
+    yield sgroup_name
     async with db_engine.begin() as conn:
         await conn.execute(
             sgroups_for_domains.delete().where(sgroups_for_domains.c.scaling_group == sgroup_name)
         )
         await conn.execute(scaling_groups.delete().where(scaling_groups.c.name == sgroup_name))
 
 
+@pytest.fixture()
+async def scaling_group_id(
+    db_engine: SAEngine,
+    scaling_group_name: ResourceGroupName,
+) -> ResourceGroupID:
+    """Return the inserted scaling group's ID."""
+    async with db_engine.begin() as conn:
+        result = await conn.execute(
+            sa.select(scaling_groups.c.id).where(scaling_groups.c.name == scaling_group_name)
+        )
+        return ResourceGroupID(result.scalar_one())
+
+
 @pytest.fixture()
 async def group_fixture(
     db_engine: SAEngine,
@@ -936,7 +946,7 @@ async def regular_user_fixture(
 async def database_fixture(
     admin_user_fixture: UserFixtureData,
     regular_user_fixture: UserFixtureData,
-    scaling_group_fixture: ScalingGroupFixtureData,
+    scaling_group_name: ResourceGroupName,
 ) -> AsyncIterator[None]:
     """Backward-compatible aggregate: requests all seed data fixtures."""
     yield