2 # Copyright (C) The Arvados Authors. All rights reserved.
4 # SPDX-License-Identifier: AGPL-3.0
6 from __future__ import absolute_import, print_function
7 from future import standard_library
15 import arvnodeman.status as status
16 import arvnodeman.config as config
19 class TestServer(object):
20 def __init__(self, management_token=None):
21 self.mgmt_token = management_token
24 cfg = config.NodeManagerConfig()
25 cfg.set('Manage', 'port', '0')
26 cfg.set('Manage', 'address', '127.0.0.1')
27 if self.mgmt_token != None:
28 cfg.set('Manage', 'ManagementToken', self.mgmt_token)
29 self.srv = status.Server(cfg)
31 addr, port = self.srv.server_address
32 self.srv_base = 'http://127.0.0.1:'+str(port)
35 def __exit__(self, exc_type, exc_value, traceback):
38 def get_status_response(self):
39 return requests.get(self.srv_base+'/status.json')
42 return self.get_status_response().json()
44 def get_healthcheck_ping(self, auth_header=None):
46 if auth_header != None:
47 headers['Authorization'] = auth_header
48 return requests.get(self.srv_base+'/_health/ping', headers=headers)
50 class StatusServerUpdates(unittest.TestCase):
51 def test_updates(self):
52 with TestServer() as srv:
54 status.tracker.update({'nodes_'+str(n): n})
55 r = srv.get_status_response()
56 self.assertEqual(200, r.status_code)
57 self.assertEqual('application/json', r.headers['content-type'])
59 self.assertEqual(n, resp['nodes_'+str(n)])
60 self.assertEqual(1, resp['nodes_1'])
61 self.assertIn('Version', resp)
62 self.assertIn('config_max_nodes', resp)
64 def test_counters(self):
65 with TestServer() as srv:
66 resp = srv.get_status()
67 # Test counters existance
68 for counter in ['list_nodes_errors', 'create_node_errors',
69 'destroy_node_errors', 'boot_failures', 'actor_exceptions']:
70 self.assertIn(counter, resp)
71 # Test counter increment
72 for count in range(1, 3):
73 status.tracker.counter_add('a_counter')
74 resp = srv.get_status()
75 self.assertEqual(count, resp['a_counter'])
77 @mock.patch('time.time')
78 def test_idle_times(self, time_mock):
79 with TestServer() as srv:
80 resp = srv.get_status()
81 node_name = 'idle_compute{}'.format(random.randint(1, 1024))
82 self.assertIn('idle_times', resp)
83 # Test add an idle node
84 time_mock.return_value = 10
85 status.tracker.idle_in(node_name)
86 time_mock.return_value += 10
87 resp = srv.get_status()
88 self.assertEqual(10, resp['idle_times'][node_name])
89 # Test adding the same idle node a 2nd time
90 time_mock.return_value += 10
91 status.tracker.idle_in(node_name)
92 time_mock.return_value += 10
93 resp = srv.get_status()
94 # Idle timestamp doesn't get reset if already exists
95 self.assertEqual(30, resp['idle_times'][node_name])
96 # Test remove idle node
97 status.tracker.idle_out(node_name)
98 resp = srv.get_status()
99 self.assertNotIn(node_name, resp['idle_times'])
102 class StatusServerDisabled(unittest.TestCase):
103 def test_config_disabled(self):
104 cfg = config.NodeManagerConfig()
105 cfg.set('Manage', 'port', '-1')
106 cfg.set('Manage', 'address', '127.0.0.1')
107 self.srv = status.Server(cfg)
109 self.assertFalse(self.srv.enabled)
110 self.assertFalse(getattr(self.srv, '_thread', False))
112 class HealthcheckPing(unittest.TestCase):
113 def test_ping_disabled(self):
114 with TestServer() as srv:
115 r = srv.get_healthcheck_ping()
116 self.assertEqual(404, r.status_code)
118 def test_ping_no_auth(self):
119 with TestServer('configuredmanagementtoken') as srv:
120 r = srv.get_healthcheck_ping()
121 self.assertEqual(401, r.status_code)
123 def test_ping_bad_auth_format(self):
124 with TestServer('configuredmanagementtoken') as srv:
125 r = srv.get_healthcheck_ping('noBearer')
126 self.assertEqual(403, r.status_code)
128 def test_ping_bad_auth_token(self):
129 with TestServer('configuredmanagementtoken') as srv:
130 r = srv.get_healthcheck_ping('Bearer badtoken')
131 self.assertEqual(403, r.status_code)
133 def test_ping_success(self):
134 with TestServer('configuredmanagementtoken') as srv:
135 r = srv.get_healthcheck_ping('Bearer configuredmanagementtoken')
136 self.assertEqual(200, r.status_code)
137 self.assertEqual('application/json', r.headers['content-type'])
139 self.assertEqual('{"health": "OK"}', json.dumps(resp))