test/scheduler: fix

m-labs · Feb 15, 2016 · 1fc51f3 · 1fc51f3
1 parent 72a993a
commit 1fc51f3
Showing 1 changed file with 13 additions and 3 deletions.
diff --git a/artiq/test/test_scheduler.py b/artiq/test/test_scheduler.py
@@ -63,6 +63,16 @@ def _get_basic_steps(rid, expid, priority=0, flush=False):
     ]
 
 
+class _RIDCounter:
+    def __init__(self, next_rid):
+        self._next_rid = next_rid
+
+    def get(self):
+        rid = self._next_rid
+        self._next_rid += 1
+        return rid
+
+
 class SchedulerCase(unittest.TestCase):
     def setUp(self):
         if os.name == "nt":
@@ -73,7 +83,7 @@ def setUp(self):
 
     def test_steps(self):
         loop = self.loop
-        scheduler = Scheduler(0, dict(), None)
+        scheduler = Scheduler(_RIDCounter(0), dict(), None)
         expid = _get_expid("EmptyExperiment")
 
         expect = _get_basic_steps(1, expid)
@@ -121,7 +131,7 @@ def check_termination(mod):
         handlers = {
             "update_dataset": check_termination
         }
-        scheduler = Scheduler(0, handlers, None)
+        scheduler = Scheduler(_RIDCounter(0), handlers, None)
 
         expid_bg = _get_expid("BackgroundExperiment")
         expid = _get_expid("EmptyExperiment")
@@ -165,7 +175,7 @@ def notify(mod):
 
     def test_flush(self):
         loop = self.loop
-        scheduler = Scheduler(0, dict(), None)
+        scheduler = Scheduler(_RIDCounter(0), dict(), None)
         expid = _get_expid("EmptyExperiment")
 
         expect = _get_basic_steps(1, expid, 1, True)