Fix CI

lablup · Oct 26, 2023 · 0c96c0e · 0c96c0e
1 parent 3f8d232
commit 0c96c0e
Showing 1 changed file with 24 additions and 6 deletions.
diff --git a/tests/manager/test_scheduler.py b/tests/manager/test_scheduler.py
@@ -49,9 +49,18 @@
 
 def test_load_intrinsic():
     default_sgroup_opts = ScalingGroupOpts()
-    assert isinstance(load_scheduler("fifo", default_sgroup_opts, {}), FIFOSlotScheduler)
-    assert isinstance(load_scheduler("lifo", default_sgroup_opts, {}), LIFOSlotScheduler)
-    assert isinstance(load_scheduler("drf", default_sgroup_opts, {}), DRFScheduler)
+    assert isinstance(
+        load_scheduler("fifo", default_sgroup_opts, {}, agent_selection_resource_priority),
+        FIFOSlotScheduler,
+    )
+    assert isinstance(
+        load_scheduler("lifo", default_sgroup_opts, {}, agent_selection_resource_priority),
+        LIFOSlotScheduler,
+    )
+    assert isinstance(
+        load_scheduler("drf", default_sgroup_opts, {}, agent_selection_resource_priority),
+        DRFScheduler,
+    )
 
 
 def test_scheduler_configs():
@@ -64,15 +73,22 @@ def test_scheduler_configs():
             "num_retries_to_skip": 5,
         },
     )
-    scheduler = load_scheduler("fifo", example_sgroup_opts, example_sgroup_opts.config)
+    scheduler = load_scheduler(
+        "fifo", example_sgroup_opts, example_sgroup_opts.config, agent_selection_resource_priority
+    )
     assert isinstance(scheduler, FIFOSlotScheduler)
     assert scheduler.config == {
         "extra_config": None,
         "num_retries_to_skip": 5,
     }
     with pytest.raises(t.DataError):
         example_sgroup_opts.config["num_retries_to_skip"] = -1  # invalid value
-        scheduler = load_scheduler("fifo", example_sgroup_opts, example_sgroup_opts.config)
+        scheduler = load_scheduler(
+            "fifo",
+            example_sgroup_opts,
+            example_sgroup_opts.config,
+            agent_selection_resource_priority,
+        )
 
 
 example_group_id = uuid4()
@@ -891,7 +907,9 @@ def test_fifo_scheduler_hol_blocking_avoidance_config():
     picked_session_id = scheduler.pick_session(example_total_capacity, pending_sessions, [])
     assert picked_session_id == "s0"
 
-    scheduler = FIFOSlotScheduler(ScalingGroupOpts(), {"num_retries_to_skip": 5})
+    scheduler = FIFOSlotScheduler(
+        ScalingGroupOpts(), {"num_retries_to_skip": 5}, agent_selection_resource_priority
+    )
     pending_sessions = [
         gen_pending_for_holb_tests("s0", {"scheduler": {"retries": 5}}),
         gen_pending_for_holb_tests("s1", {"scheduler": {"retries": 4}}),