diff --git a/artiq/test/test_scheduler.py b/artiq/test/test_scheduler.py index 7e90c78c14..95df02ec0a 100644 --- a/artiq/test/test_scheduler.py +++ b/artiq/test/test_scheduler.py @@ -84,25 +84,33 @@ def get(self): return rid class SchedulerMonitor(): - def __init__(self, expect_status): - self.expect_status = expect_status + def __init__(self): self.experiments = {} - self.current_status = {} + self.last_status = {} + self.records = {} + self.finished = False - def check_status(self, test): + def record(self): for key, value in self.experiments.items(): - if key not in self.current_status.keys(): - self.current_status[key] = "" - if self.experiments[key]["status"] != self.current_status[key]: - test.assertEqual(self.experiments[key]["status"], self.expect_status[key][0]) - self.current_status[key] = self.expect_status[key].pop(0) - - def finished(self, test): - for value in self.expect_status.values(): - if value: - return False - return True - + if key not in self.last_status.keys(): + self.last_status[key] = "" + self.records[key] = [] + current_status = self.experiments[key]["status"] + if current_status != self.last_status[key]: + self.last_status[key] = current_status + self.records[key].append(time()) + self.records[key].append(current_status) + if current_status == "deleting": + self.finished = True + + def get_in_time(self, rid, status): + return self.records[rid][self.records[rid].index(status)-1] + + def get_out_time(self, rid, status) + if self.records[rid][-1] == status: + return "never" + else: + return self.records[rid][self.records[rid].index(status) + 1] class SchedulerCase(unittest.TestCase): def setUp(self): @@ -165,23 +173,15 @@ def test_pending_priority(self): late = time() + 100000 early = time() + 1 - expect_status = { - 0: ["pending", "preparing", "prepare_done", - "running", "paused", "running"], - 1: ["pending"], - 2: ["pending", "preparing", "prepare_done", - "running", "run_done", "analyzing", "deleting"], - } - - scheduler_mon = SchedulerMonitor(expect_status) + scheduler_mon = SchedulerMonitor() done = asyncio.Event() def notify(mod): process_mod(scheduler_mon.experiments, mod) - scheduler_mon.check_status(self) + scheduler_mon.record() - if scheduler_mon.finished(self): + if scheduler_mon.finished: done.set() scheduler.notifier.publish = notify @@ -196,6 +196,10 @@ def notify(mod): scheduler.notifier.publish = None loop.run_until_complete(scheduler.stop()) + # Assert + self.assertTrue(scheduler_mon.get_time(1, "out", "pending") == "never") + self.assertTrue(scheduler_mon.get_time(2, "out", "pending") >= early) + def test_pause(self): loop = self.loop