X-Git-Url: https://git.arvados.org/arvados.git/blobdiff_plain/987c225f9f6845868ee674902090c27a5a064f42..5141c3ee23e89696773e227a93236ef2a51543c2:/services/nodemanager/tests/test_computenode.py

diff --git a/services/nodemanager/tests/test_computenode.py b/services/nodemanager/tests/test_computenode.py
index 477e20eb57..5ced5f99bb 100644
--- a/services/nodemanager/tests/test_computenode.py
+++ b/services/nodemanager/tests/test_computenode.py
@@ -114,6 +114,14 @@ class ComputeNodeShutdownActorTestCase(testutil.ActorTestMixin,
         self.stop_proxy(self.shutdown_actor)
         self.assertTrue(self.cloud_client.destroy_node.called)
 
+    def test_late_subscribe(self):
+        self.make_actor()
+        subscriber = mock.Mock(name='subscriber_mock')
+        self.shutdown_actor.subscribe(subscriber).get(self.TIMEOUT)
+        self.stop_proxy(self.shutdown_actor)
+        self.assertEqual(self.shutdown_actor.actor_ref.actor_urn,
+                         subscriber.call_args[0][0].actor_ref.actor_urn)
+
 
 class ComputeNodeUpdateActorTestCase(testutil.ActorTestMixin,
                                      unittest.TestCase):
@@ -178,11 +186,38 @@ class ComputeNodeMonitorActorTestCase(testutil.ActorTestMixin,
         self.node_actor = cnode.ComputeNodeMonitorActor.start(
             self.cloud_mock, start_time, self.shutdowns, self.timer,
             self.updates, arv_node).proxy()
-        self.subscription = self.node_actor.subscribe(self.subscriber)
+        self.node_actor.subscribe(self.subscriber).get(self.TIMEOUT)
+
+    def node_state(self, *states):
+        return self.node_actor.in_state(*states).get(self.TIMEOUT)
+
+    def test_in_state_when_unpaired(self):
+        self.make_actor()
+        self.assertIsNone(self.node_state('idle', 'alloc'))
+
+    def test_in_state_when_pairing_stale(self):
+        self.make_actor(arv_node=testutil.arvados_node_mock(
+                job_uuid=None, age=90000))
+        self.assertIsNone(self.node_state('idle', 'alloc'))
+
+    def test_in_state_when_no_state_available(self):
+        self.make_actor(arv_node=testutil.arvados_node_mock(info={}))
+        self.assertIsNone(self.node_state('idle', 'alloc'))
+
+    def test_in_idle_state(self):
+        self.make_actor(2, arv_node=testutil.arvados_node_mock(job_uuid=None))
+        self.assertTrue(self.node_state('idle'))
+        self.assertFalse(self.node_state('alloc'))
+        self.assertTrue(self.node_state('idle', 'alloc'))
+
+    def test_in_alloc_state(self):
+        self.make_actor(3, arv_node=testutil.arvados_node_mock(job_uuid=True))
+        self.assertFalse(self.node_state('idle'))
+        self.assertTrue(self.node_state('alloc'))
+        self.assertTrue(self.node_state('idle', 'alloc'))
 
     def test_init_shutdown_scheduling(self):
         self.make_actor()
-        self.subscription.get(self.TIMEOUT)
         self.assertTrue(self.timer.schedule.called)
         self.assertEqual(300, self.timer.schedule.call_args[0][0])
 
@@ -230,7 +265,7 @@ class ComputeNodeMonitorActorTestCase(testutil.ActorTestMixin,
         self.check_shutdown_rescheduled(True, 600)
 
     def test_no_shutdown_when_node_state_stale(self):
-        self.make_actor(6, testutil.arvados_node_mock(6, age=900))
+        self.make_actor(6, testutil.arvados_node_mock(6, age=90000))
         self.check_shutdown_rescheduled(True, 600)
 
     def test_arvados_node_match(self):