services/nodemanager/tests/test_nodelist.py

   1 #!/usr/bin/env python
   2 # Copyright (C) The Arvados Authors. All rights reserved.
   3 #
   4 # SPDX-License-Identifier: AGPL-3.0
   5
   6 from __future__ import absolute_import, print_function
   7
   8 import unittest
   9 import mock
  10
  11 import arvnodeman.nodelist as nodelist
  12 from libcloud.compute.base import NodeSize
  13 from . import testutil
  14
  15 class ArvadosNodeListMonitorActorTestCase(testutil.RemotePollLoopActorTestMixin,
  16                                           unittest.TestCase):
  17     TEST_CLASS = nodelist.ArvadosNodeListMonitorActor
  18
  19     def build_monitor(self, side_effect, *args, **kwargs):
  20         super(ArvadosNodeListMonitorActorTestCase, self).build_monitor(
  21             *args, **kwargs)
  22         self.client.nodes().list().execute.side_effect = side_effect
  23
  24     @mock.patch("subprocess.check_output")
  25     def test_uuid_is_subscription_key(self, sinfo_mock):
  26         sinfo_mock.return_value = ""
  27         node = testutil.arvados_node_mock()
  28         self.build_monitor([{
  29             'items': [node],
  30             'items_available': 1,
  31             'offset': 0
  32         }, {
  33             'items': [],
  34             'items_available': 1,
  35             'offset': 1
  36         }])
  37         self.monitor.subscribe_to(node['uuid'],
  38                                   self.subscriber).get(self.TIMEOUT)
  39         self.stop_proxy(self.monitor)
  40         self.subscriber.assert_called_with(node)
  41         self.assertEqual("down", node["crunch_worker_state"])
  42
  43     @mock.patch("subprocess.check_output")
  44     def test_update_from_sinfo(self, sinfo_mock):
  45         sinfo_mock.return_value = """compute1|idle|instancetype=a1.test
  46 compute2|alloc|(null)
  47 notarvados12345|idle|(null)
  48 """
  49         nodeIdle = testutil.arvados_node_mock(node_num=1)
  50         nodeBusy = testutil.arvados_node_mock(node_num=2)
  51         nodeMissing = testutil.arvados_node_mock(node_num=99)
  52         self.build_monitor([{
  53             'items': [nodeIdle, nodeBusy, nodeMissing],
  54             'items_available': 1,
  55             'offset': 0
  56         }, {
  57             'items': [],
  58             'items_available': 1,
  59             'offset': 1
  60         }])
  61         self.monitor.subscribe_to(nodeMissing['uuid'],
  62                                   self.subscriber).get(self.TIMEOUT)
  63         self.stop_proxy(self.monitor)
  64         self.subscriber.assert_called_with(nodeMissing)
  65
  66         self.assertEqual("idle", nodeIdle["crunch_worker_state"])
  67         self.assertEqual("busy", nodeBusy["crunch_worker_state"])
  68         self.assertEqual("down", nodeMissing["crunch_worker_state"])
  69
  70         self.assertEqual("instancetype=a1.test", nodeIdle["slurm_node_features"])
  71         self.assertEqual("", nodeBusy["slurm_node_features"])
  72         self.assertEqual("", nodeMissing["slurm_node_features"])
  73
  74
  75 class CloudNodeListMonitorActorTestCase(testutil.RemotePollLoopActorTestMixin,
  76                                         unittest.TestCase):
  77     TEST_CLASS = nodelist.CloudNodeListMonitorActor
  78
  79     class MockNode(object):
  80         def __init__(self, count):
  81             self.id = str(count)
  82             self.name = 'test{}.example.com'.format(count)
  83             self.private_ips = ['10.0.0.{}'.format(count)]
  84             self.public_ips = []
  85             self.size = testutil.MockSize(1)
  86             self.state = 0
  87             self.extra = {'arvados_node_size': self.size.id}
  88
  89
  90     def build_monitor(self, side_effect, *args, **kwargs):
  91         super(CloudNodeListMonitorActorTestCase, self).build_monitor(
  92             *args, **kwargs)
  93         self.client.list_nodes.side_effect = side_effect
  94
  95     def test_id_is_subscription_key(self):
  96         node = self.MockNode(1)
  97         mock_calc = mock.MagicMock()
  98         mock_calc.find_size.return_value = testutil.MockSize(2)
  99         self.build_monitor([[node]], mock_calc)
 100         self.monitor.subscribe_to('1', self.subscriber).get(self.TIMEOUT)
 101         self.stop_proxy(self.monitor)
 102         self.subscriber.assert_called_with(node)
 103         self.assertEqual(testutil.MockSize(2), node.size)
 104
 105 if __name__ == '__main__':
 106     unittest.main()