scheduler: Test RID 2 is chosen for running after RID 0 has paused

RID 2 has earlier due_date than RID 1 so it is chosen for running. RID 2 should complete all basic steps. Signed-off-by: yuk <[email protected]>
m-labs · Oct 11, 2022 · 6681586 · 6681586
1 parent 3ffbc56
commit 6681586
Showing 1 changed file with 50 additions and 120 deletions.
diff --git a/artiq/test/test_scheduler.py b/artiq/test/test_scheduler.py
@@ -144,137 +144,67 @@ def test_pending_priority(self):
         late = time() + 100000
         early = time() + 1
 
-        expect = [
-            {
-                "path": [],
-                "action": "setitem",
-                "value": {
-                    "repo_msg": None,
-                    "priority": low_priority,
-                    "pipeline": "main",
-                    "due_date": None,
-                    "status": "pending",
-                    "expid": expid_bg,
-                    "flush": False
-                },
-                "key": 0
-            },
-            {
-                "path": [],
-                "action": "setitem",
-                "value": {
-                    "repo_msg": None,
-                    "priority": high_priority,
-                    "pipeline": "main",
-                    "due_date": late,
-                    "status": "pending",
-                    "expid": expid_empty,
-                    "flush": False
-                },
-                "key": 1
-            },
-            {
-                "path": [],
-                "action": "setitem",
-                "value": {
-                    "repo_msg": None,
-                    "priority": middle_priority,
-                    "pipeline": "main",
-                    "due_date": early,
-                    "status": "pending",
-                    "expid": expid_empty,
-                    "flush": False
-                },
-                "key": 2
-            },
-            {
-                "path": [0],
-                "action": "setitem",
-                "value": "preparing",
-                "key": "status"
-            },
-            {
-                "path": [0],
-                "action": "setitem",
-                "value": "prepare_done",
-                "key": "status"
-            },
-            {
-                "path": [0],
-                "action": "setitem",
-                "value": "running",
-                "key": "status"
-            },
-            {
-                "path": [2],
-                "action": "setitem",
-                "value": "preparing",
-                "key": "status"
-            },
-            {
-                "path": [2],
-                "action": "setitem",
-                "value": "prepare_done",
-                "key": "status"
-            },
-            {
-                "path": [0],
-                "action": "setitem",
-                "value": "paused",
-                "key": "status"
-            },
-            {
-                "path": [2],
-                "action": "setitem",
-                "value": "running",
-                "key": "status"
-            },
-            {
-                "path": [2],
-                "action": "setitem",
-                "value": "run_done",
-                "key": "status"
-            },
-            {
-                "path": [0],
-                "action": "setitem",
-                "value": "running",
-                "key": "status"
-            },
-            {
-                "path": [2],
-                "action": "setitem",
-                "value": "analyzing",
-                "key": "status"
-            },
-            {
-                "path": [2],
-                "action": "setitem",
-                "value": "deleting",
-                "key": "status"
-            },
-            {
-                "path": [],
-                "action": "delitem",
-                "key": 2
-            },
-        ]
+        expect_rid2 = _get_basic_steps(2, expid_empty, middle_priority)
+        expect_rid2[0]["value"].update(due_date=early)
+        rid0_paused = asyncio.Event()
+        rid1_running = asyncio.Event()
+        rid2_running = asyncio.Event()
         done = asyncio.Event()
-        expect_idx = 0
+        expect_rid2_idx = 0
+
         def notify(mod):
-            nonlocal expect_idx
-            self.assertEqual(mod, expect[expect_idx])
-            expect_idx += 1
-            if expect_idx >= len(expect):
+            nonlocal expect_rid2_idx
+            if mod == {"path": [0],
+                       "value": "paused",
+                       "key": "status",
+                       "action": "setitem"}:
+                rid0_paused.set()
+            if mod == {"path": [1],
+                       "value": "running",
+                       "key": "status",
+                       "action": "setitem"}:
+                rid1_running.set()
+            if mod == {"path": [2],
+                       "value": "running",
+                       "key": "status",
+                       "action": "setitem"}:
+                rid2_running.set()
+            if mod["path"] == [2] or (mod["path"] == [] and mod["key"] == 2):
+                self.assertEqual(mod, expect_rid2[expect_rid2_idx])
+                expect_rid2_idx += 1
+            if expect_rid2_idx >= len(expect_rid2):
                 done.set()
         scheduler.notifier.publish = notify
 
+        async def expect_paused_running():
+            rid1_running_future = asyncio.ensure_future(rid1_running.wait())
+            rid2_running_future = asyncio.ensure_future(rid2_running.wait())
+            # expect RID 0 paused -> RID 2 running
+            await rid0_paused.wait()
+            done, pending = await asyncio.wait(
+                [rid1_running_future, rid2_running_future],
+                return_when=asyncio.FIRST_COMPLETED
+            )
+            assert rid2_running_future in done and \
+                   rid1_running_future in pending
+            for task in pending:
+                task.cancel()
+
         scheduler.start()
 
         scheduler.submit("main", expid_bg, low_priority)
         scheduler.submit("main", expid_empty, high_priority, late)
         scheduler.submit("main", expid_empty, middle_priority, early)
 
+        timeout = 5
+        try:
+            loop.run_until_complete(
+                asyncio.wait_for(expect_paused_running(), timeout)
+            )
+        except asyncio.TimeoutError:
+            raise AssertionError(
+                f"expect_paused_running() did not complete within {timeout}s"
+            )
         loop.run_until_complete(done.wait())
         scheduler.notifier.publish = None
         loop.run_until_complete(scheduler.stop())