diff options
Diffstat (limited to 'mobly/base_test.py')
| -rw-r--r-- | mobly/base_test.py | 134 |
1 files changed, 130 insertions, 4 deletions
diff --git a/mobly/base_test.py b/mobly/base_test.py index f1a28bf..ebcc2b8 100644 --- a/mobly/base_test.py +++ b/mobly/base_test.py @@ -43,11 +43,86 @@ STAGE_NAME_TEARDOWN_TEST = 'teardown_test' STAGE_NAME_TEARDOWN_CLASS = 'teardown_class' STAGE_NAME_CLEAN_UP = 'clean_up' +# Attribute names +ATTR_REPEAT_CNT = '_repeat_count' +ATTR_MAX_RETRY_CNT = '_max_count' + class Error(Exception): """Raised for exceptions that occurred in BaseTestClass.""" +def repeat(count): + """Decorator for repeating a test case multiple times. + + The BaseTestClass will execute the test cases annotated with this decorator + the specified number of time. + + This decorator only stores the information needed for the repeat. It does not + execute the repeat. + + Args: + count: int, the total number of times to execute the decorated test case. + + Returns: + The wrapped test function. + + Raises: + ValueError, if the user input is invalid. + """ + if count <= 1: + raise ValueError( + f'The `count` for `repeat` must be larger than 1, got "{count}".') + + def _outer_decorator(func): + setattr(func, ATTR_REPEAT_CNT, count) + + @functools.wraps(func) + def _wrapper(*args): + func(*args) + + return _wrapper + + return _outer_decorator + + +def retry(max_count): + """Decorator for retrying a test case until it passes. + + The BaseTestClass will keep executing the test cases annotated with this + decorator until the test passes, or the maxinum number of iterations have + been met. + + This decorator only stores the information needed for the retry. It does not + execute the retry. + + Args: + max_count: int, the maximum number of times to execute the decorated test + case. + + Returns: + The wrapped test function. + + Raises: + ValueError, if the user input is invalid. + """ + if max_count <= 1: + raise ValueError( + f'The `max_count` for `retry` must be larger than 1, got "{max_count}".' + ) + + def _outer_decorator(func): + setattr(func, ATTR_MAX_RETRY_CNT, max_count) + + @functools.wraps(func) + def _wrapper(*args): + func(*args) + + return _wrapper + + return _outer_decorator + + class BaseTestClass: """Base class for all test classes to inherit from. @@ -559,7 +634,38 @@ class BaseTestClass: content['timestamp'] = utils.get_current_epoch_time() self.summary_writer.dump(content, records.TestSummaryEntryType.USER_DATA) - def exec_one_test(self, test_name, test_method): + def _exec_one_test_with_retry(self, test_name, test_method, max_count): + """Executes one test and retry the test if needed. + + Repeatedly execute a test case until it passes or the maximum count of + iteration has been reached. + + Args: + test_name: string, Name of the test. + test_method: function, The test method to execute. + max_count: int, the maximum number of iterations to execute the test for. + """ + + def should_retry(record): + return record.result in [ + records.TestResultEnums.TEST_RESULT_FAIL, + records.TestResultEnums.TEST_RESULT_ERROR + ] + + previous_record = self.exec_one_test(test_name, test_method) + + if not should_retry(previous_record): + return previous_record + + for i in range(max_count - 1): + retry_name = f'{test_name}_retry_{i+1}' + new_record = records.TestResultRecord(retry_name, self.TAG) + new_record.retry_parent = previous_record.signature + previous_record = self.exec_one_test(retry_name, test_method, new_record) + if not should_retry(previous_record): + break + + def exec_one_test(self, test_name, test_method, record=None): """Executes one test and update test results. Executes setup_test, the test method, and teardown_test; then creates a @@ -569,8 +675,18 @@ class BaseTestClass: Args: test_name: string, Name of the test. test_method: function, The test method to execute. + record: records.TestResultRecord, optional arg for injecting a record + object to use for this test execution. If not set, a new one is created + created. This is meant for passing infomation between consecutive test + case execution for retry purposes. Do NOT abuse this for "magical" + features. + + Returns: + TestResultRecord, the test result record object of the test execution. + This object is strictly for read-only purposes. Modifying this record + will not change what is reported in the test run's summary yaml file. """ - tr_record = records.TestResultRecord(test_name, self.TAG) + tr_record = record or records.TestResultRecord(test_name, self.TAG) tr_record.uid = getattr(test_method, 'uid', None) tr_record.test_begin() self.current_test_info = runtime_test_info.RuntimeTestInfo( @@ -650,6 +766,7 @@ class BaseTestClass: self.summary_writer.dump(tr_record.to_dict(), records.TestSummaryEntryType.RECORD) self.current_test_info = None + return tr_record def _assert_function_name_in_stack(self, expected_func_name): """Asserts that the current stack contains the given function name.""" @@ -767,7 +884,12 @@ class BaseTestClass: test_method = self._generated_test_table[test_name] else: raise Error('%s does not have test method %s.' % (self.TAG, test_name)) - test_methods.append((test_name, test_method)) + repeat_count = getattr(test_method, ATTR_REPEAT_CNT, 0) + if repeat_count: + for i in range(repeat_count): + test_methods.append((f'{test_name}_{i}', test_method)) + else: + test_methods.append((test_name, test_method)) return test_methods def _skip_remaining_tests(self, exception): @@ -831,7 +953,11 @@ class BaseTestClass: return setup_class_result # Run tests in order. for test_name, test_method in tests: - self.exec_one_test(test_name, test_method) + max_count = getattr(test_method, ATTR_MAX_RETRY_CNT, 0) + if max_count: + self._exec_one_test_with_retry(test_name, test_method, max_count) + else: + self.exec_one_test(test_name, test_method) return self.results except signals.TestAbortClass as e: e.details = 'Test class aborted due to: %s' % e.details |