From 0c96c0ec6d1a8ffa39a0ffdffffa7e8013cc064a Mon Sep 17 00:00:00 2001
From: Gyubong Lee <jopemachine@naver.com>
Date: Thu, 26 Oct 2023 04:58:07 +0000
Subject: [PATCH] Fix CI

---
 tests/manager/test_scheduler.py | 30 ++++++++++++++++++++++++------
 1 file changed, 24 insertions(+), 6 deletions(-)

diff --git a/tests/manager/test_scheduler.py b/tests/manager/test_scheduler.py
index 43e90897006..d69123e0bc4 100644
--- 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,7 +73,9 @@ 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,
@@ -72,7 +83,12 @@ def test_scheduler_configs():
     }
     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}}),