12055: Merge branch 'master' into 12055-nodemanager-ec2-tags
[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("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 = "compute99 alloc"
46         node = testutil.arvados_node_mock()
47         self.build_monitor([{
48             'items': [node],
49             'items_available': 1,
50             'offset': 0
51         }, {
52             'items': [],
53             'items_available': 1,
54             'offset': 1
55         }])
56         self.monitor.subscribe_to(node['uuid'],
57                                   self.subscriber).get(self.TIMEOUT)
58         self.stop_proxy(self.monitor)
59         self.subscriber.assert_called_with(node)
60         self.assertEqual("busy", node["crunch_worker_state"])
61
62
63 class CloudNodeListMonitorActorTestCase(testutil.RemotePollLoopActorTestMixin,
64                                         unittest.TestCase):
65     TEST_CLASS = nodelist.CloudNodeListMonitorActor
66
67     class MockNode(object):
68         def __init__(self, count):
69             self.id = str(count)
70             self.name = 'test{}.example.com'.format(count)
71             self.private_ips = ['10.0.0.{}'.format(count)]
72             self.public_ips = []
73             self.size = testutil.MockSize(1)
74             self.state = 0
75
76
77     def build_monitor(self, side_effect, *args, **kwargs):
78         super(CloudNodeListMonitorActorTestCase, self).build_monitor(
79             *args, **kwargs)
80         self.client.list_nodes.side_effect = side_effect
81
82     def test_id_is_subscription_key(self):
83         node = self.MockNode(1)
84         mock_calc = mock.MagicMock()
85         mock_calc.find_size.return_value = testutil.MockSize(2)
86         self.build_monitor([[node]], mock_calc)
87         self.monitor.subscribe_to('1', self.subscriber).get(self.TIMEOUT)
88         self.stop_proxy(self.monitor)
89         self.subscriber.assert_called_with(node)
90         self.assertEqual(testutil.MockSize(2), node.size)
91
92 if __name__ == '__main__':
93     unittest.main()