testing/legion/legion_test_case.py

   1 # Copyright 2015 The Chromium Authors. All rights reserved.
   2 # Use of this source code is governed by a BSD-style license that can be
   3 # found in the LICENSE file.
   4
   5 """Adds unittest-esque functionality to Legion."""
   6
   7 import argparse
   8 import logging
   9 import sys
  10 import unittest
  11
  12 #pylint: disable=relative-import
  13 import common_lib
  14 import task_controller
  15 import task_registration_server
  16
  17 BANNER_WIDTH = 80
  18
  19
  20 class TestCase(unittest.TestCase):
  21   """Test case class with added Legion support."""
  22
  23   _registration_server = None
  24   _initialized = False
  25
  26   @classmethod
  27   def __new__(cls, *args, **kwargs):
  28     """Initialize the class and return a new instance."""
  29     cls._InitializeClass()
  30     return super(TestCase, cls).__new__(*args, **kwargs)
  31
  32   def __init__(self, test_name='runTest'):
  33     super(TestCase, self).__init__(test_name)
  34     method = getattr(self, test_name, None)
  35     if method:
  36       # Install the _RunTest method
  37       self._TestMethod = method
  38       setattr(self, test_name, self._RunTest)
  39
  40   def _RunTest(self):
  41     """Runs the test method and provides banner info and error reporting."""
  42     self._LogInfoBanner(self._testMethodName, self.shortDescription())
  43     try:
  44       return self._TestMethod()
  45     except:
  46       exc_info = sys.exc_info()
  47       logging.error('', exc_info=exc_info)
  48       raise exc_info[0], exc_info[1], exc_info[2]
  49
  50   @classmethod
  51   def _InitializeClass(cls):
  52     """Handles class level initialization.
  53
  54     There are 2 types of setup/teardown methods that always need to be run:
  55     1) Framework level setup/teardown
  56     2) Test case level setup/teardown
  57
  58     This method installs handlers in place of setUpClass and tearDownClass that
  59     will ensure both types of setup/teardown methods are called correctly.
  60     """
  61     if cls._initialized:
  62       return
  63     cls._OriginalSetUpClassMethod = cls.setUpClass
  64     cls.setUpClass = cls._HandleSetUpClass
  65     cls._OriginalTearDownClassMethod = cls.tearDownClass
  66     cls.tearDownClass = cls._HandleTearDownClass
  67     cls._initialized = True
  68
  69   @classmethod
  70   def _LogInfoBanner(cls, method_name, method_doc=None):
  71     """Formats and logs test case information."""
  72     logging.info('*' * BANNER_WIDTH)
  73     logging.info(method_name.center(BANNER_WIDTH))
  74     if method_doc:
  75       for line in method_doc.split('\n'):
  76         logging.info(line.center(BANNER_WIDTH))
  77     logging.info('*' * BANNER_WIDTH)
  78
  79   @classmethod
  80   def CreateTask(cls, *args, **kwargs):
  81     """Convenience method to create a new task."""
  82     task = task_controller.TaskController(*args, **kwargs)
  83     cls._registration_server.RegisterTaskCallback(
  84         task.otp, task.OnConnect)
  85     return task
  86
  87   @classmethod
  88   def _SetUpFramework(cls):
  89     """Perform the framework-specific setup operations."""
  90     cls._registration_server = (
  91         task_registration_server.TaskRegistrationServer())
  92     cls._registration_server.Start()
  93
  94   @classmethod
  95   def _TearDownFramework(cls):
  96     """Perform the framework-specific teardown operations."""
  97     if cls._registration_server:
  98       cls._registration_server.Shutdown()
  99     task_controller.TaskController.ReleaseAllTasks()
 100
 101   @classmethod
 102   def _HandleSetUpClass(cls):
 103     """Performs common class-level setup operations.
 104
 105     This method performs test-wide setup such as starting the registration
 106     server and then calls the original setUpClass method."""
 107     try:
 108       common_lib.InitLogging()
 109       cls._LogInfoBanner('setUpClass', 'Performs class level setup.')
 110       cls._SetUpFramework()
 111       cls._OriginalSetUpClassMethod()
 112     except:
 113       # Make sure we tear down in case of any exceptions
 114       cls._HandleTearDownClass(setup_failed=True)
 115       exc_info = sys.exc_info()
 116       logging.error('', exc_info=exc_info)
 117       raise exc_info[0], exc_info[1], exc_info[2]
 118
 119   @classmethod
 120   def _HandleTearDownClass(cls, setup_failed=False):
 121     """Performs common class-level tear down operations.
 122
 123     This method calls the original tearDownClass then performs test-wide
 124     tear down such as stopping the registration server.
 125     """
 126     cls._LogInfoBanner('tearDownClass', 'Performs class level tear down.')
 127     try:
 128       if not setup_failed:
 129         cls._OriginalTearDownClassMethod()
 130     finally:
 131       cls._TearDownFramework()
 132
 133
 134 def main():
 135   unittest.main(verbosity=0, argv=sys.argv[:1])