X-Git-Url: https://git.arvados.org/arvados.git/blobdiff_plain/6d8a27a40903f0dc61876947cecc9401edd3a32c..66c13b6055a363cb08197b8c5d040ed9a511c8ca:/services/nodemanager/tests/test_timedcallback.py

diff --git a/services/nodemanager/tests/test_timedcallback.py b/services/nodemanager/tests/test_timedcallback.py
index 60f7b81bad..1d1e6c3b11 100644
--- a/services/nodemanager/tests/test_timedcallback.py
+++ b/services/nodemanager/tests/test_timedcallback.py
@@ -14,40 +14,46 @@ from . import testutil
 @testutil.no_sleep
 class TimedCallBackActorTestCase(testutil.ActorTestMixin, unittest.TestCase):
     def test_immediate_turnaround(self):
-        future = self.FUTURE_CLASS()
+        receiver = mock.Mock()
         deliverer = timedcallback.TimedCallBackActor.start().proxy()
-        deliverer.schedule(time.time() - 1, future.set, 'immediate')
-        self.assertEqual('immediate', future.get(self.TIMEOUT))
+        deliverer.schedule(time.time() - 1, receiver,
+                           'immediate').get(self.TIMEOUT)
+        self.stop_proxy(deliverer)
+        receiver.assert_called_with('immediate')
 
     def test_delayed_turnaround(self):
-        future = self.FUTURE_CLASS()
+        receiver = mock.Mock()
         with mock.patch('time.time', return_value=0) as mock_now:
             deliverer = timedcallback.TimedCallBackActor.start().proxy()
-            deliverer.schedule(1, future.set, 'delayed')
-            self.assertRaises(pykka.Timeout, future.get, .5)
+            deliverer.schedule(1, receiver, 'delayed')
+            deliverer.schedule(3, receiver, 'failure').get(self.TIMEOUT)
+            self.assertFalse(receiver.called)
             mock_now.return_value = 2
-            self.assertEqual('delayed', future.get(self.TIMEOUT))
+            deliverer.schedule(3, receiver, 'failure').get(self.TIMEOUT)
+            self.stop_proxy(deliverer)
+        receiver.assert_called_with('delayed')
 
     def test_out_of_order_scheduling(self):
-        future1 = self.FUTURE_CLASS()
-        future2 = self.FUTURE_CLASS()
+        receiver = mock.Mock()
         with mock.patch('time.time', return_value=1.5) as mock_now:
             deliverer = timedcallback.TimedCallBackActor.start().proxy()
-            deliverer.schedule(2, future2.set, 'second')
-            deliverer.schedule(1, future1.set, 'first')
-            self.assertEqual('first', future1.get(self.TIMEOUT))
-            self.assertRaises(pykka.Timeout, future2.get, .1)
-            mock_now.return_value = 3
-            self.assertEqual('second', future2.get(self.TIMEOUT))
+            deliverer.schedule(2, receiver, 'second')
+            deliverer.schedule(1, receiver, 'first')
+            deliverer.schedule(3, receiver, 'failure').get(self.TIMEOUT)
+            receiver.assert_called_with('first')
+            mock_now.return_value = 2.5
+            deliverer.schedule(3, receiver, 'failure').get(self.TIMEOUT)
+            self.stop_proxy(deliverer)
+        receiver.assert_called_with('second')
 
     def test_dead_actors_ignored(self):
         receiver = mock.Mock(name='dead_actor', spec=pykka.ActorRef)
         receiver.tell.side_effect = pykka.ActorDeadError
         deliverer = timedcallback.TimedCallBackActor.start().proxy()
-        deliverer.schedule(time.time() - 1, receiver.tell, 'error')
-        self.wait_for_call(receiver.tell)
+        deliverer.schedule(time.time() - 1, receiver.tell,
+                           'error').get(self.TIMEOUT)
+        self.assertTrue(self.stop_proxy(deliverer), "deliverer died")
         receiver.tell.assert_called_with('error')
-        self.assertTrue(deliverer.actor_ref.is_alive(), "deliverer died")
 
 
 if __name__ == '__main__':