15964: Remove qr1hi from a few more places. Delete unused includes.
[arvados.git] / 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("subprocess32.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("subprocess32.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()