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
  60
  61 class StatusServerDisabled(unittest.TestCase):
  62     def test_config_disabled(self):
  63         cfg = config.NodeManagerConfig()
  64         cfg.set('Manage', 'port', '-1')
  65         cfg.set('Manage', 'address', '127.0.0.1')
  66         self.srv = status.Server(cfg)
  67         self.srv.start()
  68         self.assertFalse(self.srv.enabled)
  69         self.assertFalse(getattr(self.srv, '_thread', False))
  70
  71 class HealthcheckPing(unittest.TestCase):
  72     def test_ping_disabled(self):
  73         with TestServer() as srv:
  74             r = srv.get_healthcheck_ping()
  75             self.assertEqual(404, r.status_code)
  76
  77     def test_ping_no_auth(self):
  78         with TestServer('configuredmanagementtoken') as srv:
  79             r = srv.get_healthcheck_ping()
  80             self.assertEqual(401, r.status_code)
  81
  82     def test_ping_bad_auth_format(self):
  83         with TestServer('configuredmanagementtoken') as srv:
  84             r = srv.get_healthcheck_ping('noBearer')
  85             self.assertEqual(403, r.status_code)
  86
  87     def test_ping_bad_auth_token(self):
  88         with TestServer('configuredmanagementtoken') as srv:
  89             r = srv.get_healthcheck_ping('Bearer badtoken')
  90             self.assertEqual(403, r.status_code)
  91
  92     def test_ping_success(self):
  93         with TestServer('configuredmanagementtoken') as srv:
  94             r = srv.get_healthcheck_ping('Bearer configuredmanagementtoken')
  95             self.assertEqual(200, r.status_code)
  96             self.assertEqual('application/json', r.headers['content-type'])
  97             resp = r.json()
  98             self.assertEqual('{"health": "OK"}', json.dumps(resp))