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     self._output_dir = None
  40
  41   @property
  42   def output_dir(self):
  43     if not self._output_dir:
  44       self._output_dir = self.rpc.GetOutputDir()
  45     return self._output_dir
  46
  47   def _RunTest(self):
  48     """Runs the test method and provides banner info and error reporting."""
  49     self._LogInfoBanner(self._testMethodName, self.shortDescription())
  50     try:
  51       return self._TestMethod()
  52     except:
  53       exc_info = sys.exc_info()
  54       logging.error('', exc_info=exc_info)
  55       raise exc_info[0], exc_info[1], exc_info[2]
  56
  57   @classmethod
  58   def _InitializeClass(cls):
  59     """Handles class level initialization.
  60
  61     There are 2 types of setup/teardown methods that always need to be run:
  62     1) Framework level setup/teardown
  63     2) Test case level setup/teardown
  64
  65     This method installs handlers in place of setUpClass and tearDownClass that
  66     will ensure both types of setup/teardown methods are called correctly.
  67     """
  68     if cls._initialized:
  69       return
  70     cls._OriginalSetUpClassMethod = cls.setUpClass
  71     cls.setUpClass = cls._HandleSetUpClass
  72     cls._OriginalTearDownClassMethod = cls.tearDownClass
  73     cls.tearDownClass = cls._HandleTearDownClass
  74     cls._initialized = True
  75
  76   @classmethod
  77   def _LogInfoBanner(cls, method_name, method_doc=None):
  78     """Formats and logs test case information."""
  79     logging.info('*' * BANNER_WIDTH)
  80     logging.info(method_name.center(BANNER_WIDTH))
  81     if method_doc:
  82       for line in method_doc.split('\n'):
  83         logging.info(line.center(BANNER_WIDTH))
  84     logging.info('*' * BANNER_WIDTH)
  85
  86   @classmethod
  87   def CreateTask(cls, *args, **kwargs):
  88     """Convenience method to create a new task."""
  89     task = task_controller.TaskController(*args, **kwargs)
  90     cls._registration_server.RegisterTaskCallback(
  91         task.otp, task.OnConnect)
  92     return task
  93
  94   @classmethod
  95   def _SetUpFramework(cls):
  96     """Perform the framework-specific setup operations."""
  97     cls._registration_server = (
  98         task_registration_server.TaskRegistrationServer())
  99     cls._registration_server.Start()
 100
 101   @classmethod
 102   def _TearDownFramework(cls):
 103     """Perform the framework-specific teardown operations."""
 104     if cls._registration_server:
 105       cls._registration_server.Shutdown()
 106     task_controller.TaskController.ReleaseAllTasks()
 107
 108   @classmethod
 109   def _HandleSetUpClass(cls):
 110     """Performs common class-level setup operations.
 111
 112     This method performs test-wide setup such as starting the registration
 113     server and then calls the original setUpClass method."""
 114     try:
 115       common_lib.InitLogging()
 116       cls._LogInfoBanner('setUpClass', 'Performs class level setup.')
 117       cls._SetUpFramework()
 118       cls._OriginalSetUpClassMethod()
 119     except:
 120       # Make sure we tear down in case of any exceptions
 121       cls._HandleTearDownClass(setup_failed=True)
 122       exc_info = sys.exc_info()
 123       logging.error('', exc_info=exc_info)
 124       raise exc_info[0], exc_info[1], exc_info[2]
 125
 126   @classmethod
 127   def _HandleTearDownClass(cls, setup_failed=False):
 128     """Performs common class-level tear down operations.
 129
 130     This method calls the original tearDownClass then performs test-wide
 131     tear down such as stopping the registration server.
 132     """
 133     cls._LogInfoBanner('tearDownClass', 'Performs class level tear down.')
 134     try:
 135       if not setup_failed:
 136         cls._OriginalTearDownClassMethod()
 137     finally:
 138       cls._TearDownFramework()
 139
 140
 141 def main():
 142   unittest.main(verbosity=0, argv=sys.argv[:1])