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))