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 |