src/third_party/chromite/buildbot/lab_status_unittest.py

   1 #!/usr/bin/python
   2
   3 # Copyright (c) 2013 The Chromium OS Authors. All rights reserved.
   4 # Use of this source code is governed by a BSD-style license that can be
   5 # found in the LICENSE file.
   6
   7 """Unittests for lab status."""
   8
   9 import constants
  10 from mock import Mock
  11 import sys
  12 import time
  13 import urllib
  14
  15 sys.path.insert(0, constants.SOURCE_ROOT)
  16 from chromite.buildbot import lab_status
  17 from chromite.lib import cros_test_lib
  18
  19
  20 class TestLabStatus(cros_test_lib.MockTestCase):
  21   """Class that tests GetLabStatus and CheckLabStatus."""
  22
  23   def setUp(self):
  24     self.PatchObject(time, 'sleep')
  25     self.PatchObject(urllib, 'urlopen')
  26
  27   def _LabStatusFile(self, message, general_state):
  28     """Returns a file-like object with the status message written in it."""
  29     my_response = Mock()
  30     my_response.json = '{"message": "%s", "general_state": "%s"}' % (
  31         message, general_state)
  32     return my_response
  33
  34   def _TestGetLabStatusHelper(self, lab_message, general_state, expected_return,
  35                               max_attempts=5, failed_attempts=0):
  36     """Tests whether we get correct lab status.
  37
  38     Args:
  39       lab_message: A message describing lab status and
  40                    disabled boards, e.g. "Lab is Up [stumpy, kiev]"
  41       general_state: Current lab state, e.g. 'open'.
  42       expected_return: The expected return of GetLabStatus,
  43                        e.g. {'lab_is_up': True, 'message': 'Lab is up'}.
  44       max_attempts: Max attempts GetLabStatus will make to get lab status.
  45       failed_attempts: Number of failed attempts we want to mock.
  46     """
  47     return_status = self._LabStatusFile(lab_message, general_state)
  48     urlopen_side_effect = [500 for _ in range(failed_attempts)]
  49     if failed_attempts < max_attempts:
  50       urlopen_side_effect.append(200)
  51       call_count = failed_attempts + 1
  52     else:
  53       call_count = max_attempts
  54     return_status.getcode.side_effect = urlopen_side_effect
  55     urllib.urlopen.return_value = return_status
  56     return_status.read.return_value = return_status.json
  57     self.assertEqual(lab_status.GetLabStatus(max_attempts), expected_return)
  58     # pylint: disable=E1101
  59     self.assertEqual(urllib.urlopen.call_count, call_count)
  60
  61   def testGetLabStatusWithOneAttempt(self):
  62     """Tests that GetLabStatus succeeds with one attempt."""
  63     expected_return = {'lab_is_up': True, 'message': 'Lab is up'}
  64     self._TestGetLabStatusHelper('Lab is up', 'open', expected_return)
  65
  66   def testGetLabStatusWithMultipleAttempts(self):
  67     """Tests that GetLabStatus succeeds after multiple tries."""
  68     expected_return = {'lab_is_up': True, 'message': 'Lab is up'}
  69     self._TestGetLabStatusHelper('Lab is up', 'open', expected_return,
  70                               max_attempts=5, failed_attempts=3)
  71
  72   def testGetLabStatusFailsWithMultipleAttempts(self):
  73     """Tests that GetLabStatus fails after multiple tries."""
  74     expected_return = {'lab_is_up': True, 'message': ''}
  75     self._TestGetLabStatusHelper('dummy_msg', 'dummy_state', expected_return,
  76                               max_attempts=5, failed_attempts=5)
  77
  78   def _BuildLabStatus(self, lab_is_up, disabled_boards=None):
  79     """Build a dictionary representing the lab status."""
  80     if lab_is_up:
  81       message = 'Lab is Open'
  82       if disabled_boards:
  83         message += '[%s]' % (', '.join(disabled_boards))
  84     else:
  85       message = 'Lab is Down (For some reason it is down.)'
  86
  87     status = {'lab_is_up': lab_is_up, 'message': message}
  88     return status
  89
  90   def _TestCheckLabStatusHelper(self, lab_is_up):
  91     """Tests CheckLabStatus runs properly."""
  92     self.PatchObject(lab_status, 'GetLabStatus')
  93     status = self._BuildLabStatus(lab_is_up)
  94     lab_status.GetLabStatus.return_value = status
  95     if lab_is_up:
  96       lab_status.CheckLabStatus()
  97     else:
  98       self.assertRaises(lab_status.LabIsDownException,
  99                         lab_status.CheckLabStatus)
 100
 101   def testCheckLabStatusWhenLabUp(self):
 102     """Tests CheckLabStatus runs properly when lab is up."""
 103     self._TestCheckLabStatusHelper(True)
 104
 105   def testCheckLabStatusWhenLabDown(self):
 106     """Tests CheckLabStatus runs properly when lab is down."""
 107     self._TestCheckLabStatusHelper(False)
 108
 109   def testCheckLabStatusWhenBoardsDisabled(self):
 110     """Tests CheckLabStatus runs properly when some boards are disabled."""
 111     self.PatchObject(lab_status, 'GetLabStatus')
 112     status = self._BuildLabStatus(True, ['stumpy', 'kiev', 'x85-alex'])
 113     lab_status.GetLabStatus.return_value = status
 114     lab_status.CheckLabStatus('lumpy')
 115     self.assertRaises(lab_status.BoardIsDisabledException,
 116                       lab_status.CheckLabStatus,
 117                       'stumpy')
 118
 119
 120 if __name__ == '__main__':
 121   cros_test_lib.main()