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 requests
  10 import unittest
  11
  12 import arvnodeman.status as status
  13 import arvnodeman.config as config
  14
  15
  16 class TestServer(object):
  17     def __enter__(self):
  18         cfg = config.NodeManagerConfig()
  19         cfg.set('Manage', 'port', '0')
  20         cfg.set('Manage', 'address', '127.0.0.1')
  21         self.srv = status.Server(cfg)
  22         self.srv.start()
  23         addr, port = self.srv.server_address
  24         self.srv_base = 'http://127.0.0.1:'+str(port)
  25         return self
  26
  27     def __exit__(self, exc_type, exc_value, traceback):
  28         self.srv.shutdown()
  29
  30     def get_status_response(self):
  31         return requests.get(self.srv_base+'/status.json')
  32
  33     def get_status(self):
  34         return self.get_status_response().json()
  35
  36
  37 class StatusServerUpdates(unittest.TestCase):
  38     def test_updates(self):
  39         with TestServer() as srv:
  40             for n in [1, 2, 3]:
  41                 status.tracker.update({'nodes_'+str(n): n})
  42                 r = srv.get_status_response()
  43                 self.assertEqual(200, r.status_code)
  44                 self.assertEqual('application/json', r.headers['content-type'])
  45                 resp = r.json()
  46                 self.assertEqual(n, resp['nodes_'+str(n)])
  47             self.assertEqual(1, resp['nodes_1'])
  48
  49
  50 class StatusServerDisabled(unittest.TestCase):
  51     def test_config_disabled(self):
  52         cfg = config.NodeManagerConfig()
  53         cfg.set('Manage', 'port', '-1')
  54         cfg.set('Manage', 'address', '127.0.0.1')
  55         self.srv = status.Server(cfg)
  56         self.srv.start()
  57         self.assertFalse(self.srv.enabled)
  58         self.assertFalse(getattr(self.srv, '_thread', False))