| OLD | NEW |
| (Empty) |
| 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]) | |
| OLD | NEW |