15964: Remove qr1hi from a few more places. Delete unused includes.
[arvados.git] / services / nodemanager / tests / test_status.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 from future import standard_library
8
9 import json
10 import mock
11 import random
12 import requests
13 import unittest
14
15 import arvnodeman.status as status
16 import arvnodeman.config as config
17
18
19 class TestServer(object):
20     def __init__(self, management_token=None):
21         self.mgmt_token = management_token
22
23     def __enter__(self):
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)
30         self.srv.start()
31         addr, port = self.srv.server_address
32         self.srv_base = 'http://127.0.0.1:'+str(port)
33         return self
34
35     def __exit__(self, exc_type, exc_value, traceback):
36         self.srv.shutdown()
37
38     def get_status_response(self):
39         return requests.get(self.srv_base+'/status.json')
40
41     def get_status(self):
42         return self.get_status_response().json()
43
44     def get_healthcheck_ping(self, auth_header=None):
45         headers = {}
46         if auth_header != None:
47             headers['Authorization'] = auth_header
48         return requests.get(self.srv_base+'/_health/ping', headers=headers)
49
50 class StatusServerUpdates(unittest.TestCase):
51     def test_updates(self):
52         with TestServer() as srv:
53             for n in [1, 2, 3]:
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'])
58                 resp = r.json()
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)
63
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'])
76
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'])
100
101
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)
108         self.srv.start()
109         self.assertFalse(self.srv.enabled)
110         self.assertFalse(getattr(self.srv, '_thread', False))
111
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)
117
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)
122
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)
127
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)
132
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'])
138             resp = r.json()
139             self.assertEqual('{"health": "OK"}', json.dumps(resp))