fix: some glitches in tests

Author: achimnol

src/ai/backend/manager/sokovan/scheduler/selectors/utils.py

@@ -62,11 +62,14 @@ def order_slots_by_priority(
         An ordered list of slot names
     """
 
-    def get_slot_index(slot_name: str | SlotName, order_reference: list[str | SlotName]) -> int:
+    def get_slot_index(
+        slot_name: str | SlotName, order_reference: list[str | SlotName]
+    ) -> tuple[int, str]:
         try:
-            return order_reference.index(slot_name)
+            return (order_reference.index(slot_name), str(slot_name))
         except ValueError:
-            return sys.maxsize
+            # Slots missing in the priority order will be pushed back and sorted alphabetically there.
+            return (sys.maxsize, str(slot_name))
 
     # Fill missing concrete slot names, by inserting slot names after corresponding device names.
     priority_order_including_slot_names = [*priority_order]

tests/manager/conftest.py

@@ -47,6 +47,7 @@
 from ai.backend.common.lock import FileLock
 from ai.backend.common.plugin.hook import HookPluginContext
 from ai.backend.common.typed_validators import HostPortPair as HostPortPairModel
+from ai.backend.common.types import ResourceSlot
 from ai.backend.logging import LocalLogger, LogLevel
 from ai.backend.logging.config import ConsoleConfig, LogDriver, LoggingConfig
 from ai.backend.logging.types import LogFormat
@@ -566,8 +567,8 @@ async def database_fixture(
     extra_fixture_file_path = Path(extra_fixture_file.name)
 
     def fixture_json_encoder(obj: Any):
-        if isinstance(obj, Mapping):
-            return dict(obj)
+        if isinstance(obj, ResourceSlot):
+            return obj.to_json()
         if isinstance(obj, uuid.UUID):
             return str(obj)
         if isinstance(obj, datetime):
@@ -576,7 +577,8 @@ def fixture_json_encoder(obj: Any):
             return obj.value
         if isinstance(obj, yarl.URL):
             return str(obj)
-
+        if isinstance(obj, Mapping):
+            return dict(obj)
         raise TypeError(f'Fixture type "{type(obj)}" not serializable')
 
     with open(extra_fixture_file_path, "w") as f:

tests/manager/services/session/actions/test_create_cluster.py

@@ -84,7 +84,7 @@ async def test_create_cluster_with_template(
     mock_create_cluster_rpc,
     processors: SessionProcessors,
     session_repository,
-):
+) -> None:
     # Create the action using template
     action = CreateClusterAction(
         session_name=cast(str, SESSION_FIXTURE_DATA.name),
@@ -168,7 +168,7 @@ async def test_create_cluster_with_gpu_template(
     mock_create_cluster_rpc,
     processors: SessionProcessors,
     session_repository,
-):
+) -> None:
     # Create the action using GPU template
     action = CreateClusterAction(
         session_name="gpu_test_session",

tests/manager/sokovan/scheduler/selectors/test_utils.py

@@ -14,14 +14,14 @@
 class TestUtilityFunctions:
     """Test utility functions used by selectors."""
 
-    def test_count_unutilized_capabilities_no_unused(self):
+    def test_count_unutilized_capabilities_no_unused(self) -> None:
         """Test counting when all requested resources are used."""
         agent = create_agent_info(
-            available_slots={
+            available_slots=ResourceSlot({
                 "cpu": Decimal("8"),
                 "mem": Decimal("16384"),
                 "cuda.shares": Decimal("4"),
-            }
+            })
         )
         requested_slots = ResourceSlot({
             "cpu": Decimal("2"),
@@ -32,15 +32,15 @@ def test_count_unutilized_capabilities_no_unused(self):
         count = count_unutilized_capabilities(agent, requested_slots)
         assert count == 0
 
-    def test_count_unutilized_capabilities_with_unused(self):
+    def test_count_unutilized_capabilities_with_unused(self) -> None:
         """Test counting with some zero-requested resources."""
         agent = create_agent_info(
-            available_slots={
+            available_slots=ResourceSlot({
                 "cpu": Decimal("8"),
                 "mem": Decimal("16384"),
                 "cuda.shares": Decimal("4"),
                 "tpu": Decimal("2"),
-            }
+            })
         )
         requested_slots = ResourceSlot({
             "cpu": Decimal("2"),
@@ -52,19 +52,19 @@ def test_count_unutilized_capabilities_with_unused(self):
         count = count_unutilized_capabilities(agent, requested_slots)
         assert count == 2  # cuda.shares and tpu are unutilized
 
-    def test_count_unutilized_capabilities_unavailable_resources(self):
+    def test_count_unutilized_capabilities_unavailable_resources(self) -> None:
         """Test that unavailable resources are not counted."""
         agent = create_agent_info(
-            available_slots={
+            available_slots=ResourceSlot({
                 "cpu": Decimal("8"),
                 "mem": Decimal("16384"),
                 "cuda.shares": Decimal("4"),
-            },
-            occupied_slots={
+            }),
+            occupied_slots=ResourceSlot({
                 "cpu": Decimal("0"),
                 "mem": Decimal("0"),
                 "cuda.shares": Decimal("4"),  # Fully occupied
-            },
+            }),
         )
         requested_slots = ResourceSlot({
             "cpu": Decimal("2"),
@@ -75,7 +75,7 @@ def test_count_unutilized_capabilities_unavailable_resources(self):
         count = count_unutilized_capabilities(agent, requested_slots)
         assert count == 0  # cuda.shares is zero-requested but not available
 
-    def test_order_slots_by_priority_basic(self):
+    def test_order_slots_by_priority_basic(self) -> None:
         """Test basic slot ordering by priority."""
         requested_slots = ResourceSlot({
             "cpu": Decimal("1"),
@@ -88,7 +88,7 @@ def test_order_slots_by_priority_basic(self):
         result = order_slots_by_priority(requested_slots, priority_order)
         assert result == ["mem", "cpu", "cuda.shares", "disk"]
 
-    def test_order_slots_by_priority_missing_priorities(self):
+    def test_order_slots_by_priority_missing_priorities(self) -> None:
         """Test ordering when some slots are not in priority list."""
         requested_slots = ResourceSlot({
             "cpu": Decimal("1"),
@@ -104,19 +104,19 @@ def test_order_slots_by_priority_missing_priorities(self):
         # Non-priority slots should be sorted alphabetically
         assert result[1:] == sorted(result[1:])
 
-    def test_order_slots_by_priority_empty_priority(self):
+    def test_order_slots_by_priority_empty_priority(self) -> None:
         """Test ordering with empty priority list."""
         requested_slots = ResourceSlot({
             "zebra": Decimal("1"),
             "alpha": Decimal("2"),
             "beta": Decimal("3"),
         })
-        priority_order = []
+        priority_order: list[str] = []
 
         result = order_slots_by_priority(requested_slots, priority_order)
         assert result == ["alpha", "beta", "zebra"]  # Alphabetical order
 
-    def test_order_slots_by_priority_nonexistent_priorities(self):
+    def test_order_slots_by_priority_nonexistent_priorities(self) -> None:
         """Test that non-existent priority slots are ignored."""
         requested_slots = ResourceSlot({
             "cpu": Decimal("1"),
@@ -127,17 +127,17 @@ def test_order_slots_by_priority_nonexistent_priorities(self):
         result = order_slots_by_priority(requested_slots, priority_order)
         assert result == ["cpu", "mem"]  # Only requested slots appear
 
-    def test_agent_info_calculations(self):
+    def test_agent_info_calculations(self) -> None:
         """Test that AgentInfo correctly calculates available resources."""
         agent = create_agent_info(
-            available_slots={
+            available_slots=ResourceSlot({
                 "cpu": Decimal("16"),
                 "mem": Decimal("32768"),
-            },
-            occupied_slots={
+            }),
+            occupied_slots=ResourceSlot({
                 "cpu": Decimal("10"),
                 "mem": Decimal("20480"),
-            },
+            }),
         )
 
         free_slots = agent.available_slots - agent.occupied_slots

tests/manager/test_registry.py

@@ -182,7 +182,7 @@ async def test_convert_resource_spec_to_resource_slot(
         },
     }
     converted_allocations = registry.convert_resource_spec_to_resource_slot(allocations)
-    assert converted_allocations["cuda.shares"] == "4.5"
+    assert converted_allocations["cuda.shares"] == Decimal("4.5")
     allocations = {
         "cpu": {
             SlotName("cpu"): {
@@ -198,5 +198,5 @@ async def test_convert_resource_spec_to_resource_slot(
         },
     }
     converted_allocations = registry.convert_resource_spec_to_resource_slot(allocations)
-    assert converted_allocations["cpu"] == "4"
-    assert converted_allocations["ram"] == str(Decimal(BinarySize.from_str("1g")) * 3)
+    assert converted_allocations["cpu"] == Decimal("4")
+    assert converted_allocations["ram"] == Decimal(BinarySize.from_str("1g")) * 3