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 requests
  11 import unittest
  12
  13 import arvnodeman.status as status
  14 import arvnodeman.config as config
  15
  16
  17 class TestServer(object):
  18     def __init__(self, management_token=None):
  19         self.mgmt_token = management_token
  20
  21     def __enter__(self):
  22         cfg = config.NodeManagerConfig()
  23         cfg.set('Manage', 'port', '0')
  24         cfg.set('Manage', 'address', '127.0.0.1')
  25         if self.mgmt_token != None:
  26             cfg.set('Manage', 'ManagementToken', self.mgmt_token)
  27         self.srv = status.Server(cfg)
  28         self.srv.start()
  29         addr, port = self.srv.server_address
  30         self.srv_base = 'http://127.0.0.1:'+str(port)
  31         return self
  32
  33     def __exit__(self, exc_type, exc_value, traceback):
  34         self.srv.shutdown()
  35
  36     def get_status_response(self):
  37         return requests.get(self.srv_base+'/status.json')
  38
  39     def get_status(self):
  40         return self.get_status_response().json()
  41
  42     def get_healthcheck_ping(self, auth_header=None):
  43         headers = {}
  44         if auth_header != None:
  45             headers['Authorization'] = auth_header
  46         return requests.get(self.srv_base+'/_health/ping', headers=headers)
  47
  48 class StatusServerUpdates(unittest.TestCase):
  49     def test_updates(self):
  50         with TestServer() as srv:
  51             for n in [1, 2, 3]:
  52                 status.tracker.update({'nodes_'+str(n): n})
  53                 r = srv.get_status_response()
  54                 self.assertEqual(200, r.status_code)
  55                 self.assertEqual('application/json', r.headers['content-type'])
  56                 resp = r.json()
  57                 self.assertEqual(n, resp['nodes_'+str(n)])
  58             self.assertEqual(1, resp['nodes_1'])
  59             self.assertIn('Version', resp)
  60
  61
  62 class StatusServerDisabled(unittest.TestCase):
  63     def test_config_disabled(self):
  64         cfg = config.NodeManagerConfig()
  65         cfg.set('Manage', 'port', '-1')
  66         cfg.set('Manage', 'address', '127.0.0.1')
  67         self.srv = status.Server(cfg)
  68         self.srv.start()
  69         self.assertFalse(self.srv.enabled)
  70         self.assertFalse(getattr(self.srv, '_thread', False))
  71
  72 class HealthcheckPing(unittest.TestCase):
  73     def test_ping_disabled(self):
  74         with TestServer() as srv:
  75             r = srv.get_healthcheck_ping()
  76             self.assertEqual(404, r.status_code)
  77
  78     def test_ping_no_auth(self):
  79         with TestServer('configuredmanagementtoken') as srv:
  80             r = srv.get_healthcheck_ping()
  81             self.assertEqual(401, r.status_code)
  82
  83     def test_ping_bad_auth_format(self):
  84         with TestServer('configuredmanagementtoken') as srv:
  85             r = srv.get_healthcheck_ping('noBearer')
  86             self.assertEqual(403, r.status_code)
  87
  88     def test_ping_bad_auth_token(self):
  89         with TestServer('configuredmanagementtoken') as srv:
  90             r = srv.get_healthcheck_ping('Bearer badtoken')
  91             self.assertEqual(403, r.status_code)
  92
  93     def test_ping_success(self):
  94         with TestServer('configuredmanagementtoken') as srv:
  95             r = srv.get_healthcheck_ping('Bearer configuredmanagementtoken')
  96             self.assertEqual(200, r.status_code)
  97             self.assertEqual('application/json', r.headers['content-type'])
  98             resp = r.json()
  99             self.assertEqual('{"health": "OK"}', json.dumps(resp))