X-Git-Url: https://git.arvados.org/arvados.git/blobdiff_plain/dd59a50f9f3933c359930806516b43899a8b4957..8a41cc44ee196c9347785baa476a370abe77c75c:/services/nodemanager/tests/test_nodelist.py

diff --git a/services/nodemanager/tests/test_nodelist.py b/services/nodemanager/tests/test_nodelist.py
index d9f47e2605..df31a12267 100644
--- a/services/nodemanager/tests/test_nodelist.py
+++ b/services/nodemanager/tests/test_nodelist.py
@@ -1,10 +1,15 @@
 #!/usr/bin/env python
+# Copyright (C) The Arvados Authors. All rights reserved.
+#
+# SPDX-License-Identifier: AGPL-3.0
 
 from __future__ import absolute_import, print_function
 
 import unittest
+import mock
 
 import arvnodeman.nodelist as nodelist
+from libcloud.compute.base import NodeSize
 from . import testutil
 
 class ArvadosNodeListMonitorActorTestCase(testutil.RemotePollLoopActorTestMixin,
@@ -16,12 +21,55 @@ class ArvadosNodeListMonitorActorTestCase(testutil.RemotePollLoopActorTestMixin,
             *args, **kwargs)
         self.client.nodes().list().execute.side_effect = side_effect
 
-    def test_uuid_is_subscription_key(self):
+    @mock.patch("subprocess32.check_output")
+    def test_uuid_is_subscription_key(self, sinfo_mock):
+        sinfo_mock.return_value = ""
         node = testutil.arvados_node_mock()
-        self.build_monitor([{'items': [node]}])
-        self.monitor.subscribe_to(node['uuid'], self.subscriber)
-        self.wait_for_call(self.subscriber)
+        self.build_monitor([{
+            'items': [node],
+            'items_available': 1,
+            'offset': 0
+        }, {
+            'items': [],
+            'items_available': 1,
+            'offset': 1
+        }])
+        self.monitor.subscribe_to(node['uuid'],
+                                  self.subscriber).get(self.TIMEOUT)
+        self.stop_proxy(self.monitor)
         self.subscriber.assert_called_with(node)
+        self.assertEqual("down", node["crunch_worker_state"])
+
+    @mock.patch("subprocess32.check_output")
+    def test_update_from_sinfo(self, sinfo_mock):
+        sinfo_mock.return_value = """compute1|idle|instancetype=a1.test
+compute2|alloc|(null)
+notarvados12345|idle|(null)
+"""
+        nodeIdle = testutil.arvados_node_mock(node_num=1)
+        nodeBusy = testutil.arvados_node_mock(node_num=2)
+        nodeMissing = testutil.arvados_node_mock(node_num=99)
+        self.build_monitor([{
+            'items': [nodeIdle, nodeBusy, nodeMissing],
+            'items_available': 1,
+            'offset': 0
+        }, {
+            'items': [],
+            'items_available': 1,
+            'offset': 1
+        }])
+        self.monitor.subscribe_to(nodeMissing['uuid'],
+                                  self.subscriber).get(self.TIMEOUT)
+        self.stop_proxy(self.monitor)
+        self.subscriber.assert_called_with(nodeMissing)
+
+        self.assertEqual("idle", nodeIdle["crunch_worker_state"])
+        self.assertEqual("busy", nodeBusy["crunch_worker_state"])
+        self.assertEqual("down", nodeMissing["crunch_worker_state"])
+
+        self.assertEqual("instancetype=a1.test", nodeIdle["slurm_node_features"])
+        self.assertEqual("", nodeBusy["slurm_node_features"])
+        self.assertEqual("", nodeMissing["slurm_node_features"])
 
 
 class CloudNodeListMonitorActorTestCase(testutil.RemotePollLoopActorTestMixin,
@@ -34,8 +82,9 @@ class CloudNodeListMonitorActorTestCase(testutil.RemotePollLoopActorTestMixin,
             self.name = 'test{}.example.com'.format(count)
             self.private_ips = ['10.0.0.{}'.format(count)]
             self.public_ips = []
-            self.size = None
+            self.size = testutil.MockSize(1)
             self.state = 0
+            self.extra = {'arvados_node_size': self.size.id}
 
 
     def build_monitor(self, side_effect, *args, **kwargs):
@@ -45,12 +94,13 @@ class CloudNodeListMonitorActorTestCase(testutil.RemotePollLoopActorTestMixin,
 
     def test_id_is_subscription_key(self):
         node = self.MockNode(1)
-        self.build_monitor([[node]])
-        self.monitor.subscribe_to('1', self.subscriber)
-        self.wait_for_call(self.subscriber)
+        mock_calc = mock.MagicMock()
+        mock_calc.find_size.return_value = testutil.MockSize(2)
+        self.build_monitor([[node]], mock_calc)
+        self.monitor.subscribe_to('1', self.subscriber).get(self.TIMEOUT)
+        self.stop_proxy(self.monitor)
         self.subscriber.assert_called_with(node)
-
+        self.assertEqual(testutil.MockSize(2), node.size)
 
 if __name__ == '__main__':
     unittest.main()
-