1# Lint as: python3 2# 3# Copyright 2020, The Android Open Source Project 4# 5# Licensed under the Apache License, Version 2.0 (the "License"); 6# you may not use this file except in compliance with the License. 7# You may obtain a copy of the License at 8# 9# http://www.apache.org/licenses/LICENSE-2.0 10# 11# Unless required by applicable law or agreed to in writing, software 12# distributed under the License is distributed on an "AS IS" BASIS, 13# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 14# See the License for the specific language governing permissions and 15# limitations under the License. 16"""CSuite-specific automated unit test functionality.""" 17 18import argparse 19import sys 20from typing import Any, Sequence, Text 21import unittest 22 23# Export the TestCase class to reduce the number of imports tests have to list. 24TestCase = unittest.TestCase 25 26_DEVICE_SERIAL = None 27 28 29def get_device_serial() -> Text: 30 """Returns the serial of the connected device.""" 31 if not _DEVICE_SERIAL: 32 raise RuntimeError( 33 'Device serial is unset, did you call main in your test?') 34 return _DEVICE_SERIAL 35 36 37def create_arg_parser(add_help: bool = False) -> argparse.ArgumentParser: 38 """Creates a new parser that can handle the default command-line flags. 39 40 The object returned by this function can be used by other modules that want to 41 add their own command-line flags. The returned parser is intended to be passed 42 to the 'parents' argument of ArgumentParser and extend the set of default 43 flags with additional ones. 44 45 Args: 46 add_help: whether to add an option which simply displays the parser’s help 47 message; this is typically false when used from other modules that want to 48 use the returned parser as a parent argument parser. 49 50 Returns: 51 A new arg parser that can handle the default flags expected by this module. 52 """ 53 54 # The below flags are passed in by the TF Python test runner. 55 parser = argparse.ArgumentParser(add_help=add_help) 56 57 parser.add_argument('-s', '--serial', help='the device serial') 58 parser.add_argument( 59 '--test-output-file', 60 help='the file in which to store the test results', 61 required=True) 62 63 return parser 64 65 66def run_tests(args: Any, unittest_argv: Sequence[Text]) -> None: 67 """Executes a set of Python unit tests. 68 69 This function is typically used by modules that extend command-line flags. 70 Callers create their own argument parser with this module's parser as a parent 71 and parse the command-line. The resulting object is will contain the 72 attributes expected by this module and is used to call this method. 73 74 Args: 75 args: an object that contains at least the set of attributes defined in 76 objects returned when using the default argument parser. 77 unittest_argv: the list of command-line arguments to forward to 78 unittest.main. 79 """ 80 global _DEVICE_SERIAL 81 82 _DEVICE_SERIAL = args.serial 83 84 with open(args.test_output_file, 'w') as test_output_file: 85 86 # Note that we use a type and not an instance for 'testRunner' since 87 # TestProgram forwards its constructor arguments when creating an instance 88 # of the runner type. Not doing so would require us to make sure that the 89 # parameters passed to TestProgram are aligned with those for creating a 90 # runner instance. 91 class TestRunner(unittest.TextTestRunner): 92 """A test runner that writes test results to the TF-provided file.""" 93 94 def __init__(self, *args, **kwargs): 95 super(TestRunner, self).__init__( 96 stream=test_output_file, *args, **kwargs) 97 98 # Setting verbosity is required to generate output that the TradeFed test 99 # runner can parse. 100 unittest.TestProgram(verbosity=3, testRunner=TestRunner, argv=unittest_argv) 101 102 103def main(): 104 """Executes a set of Python unit tests.""" 105 args, unittest_argv = create_arg_parser(add_help=True).parse_known_args( 106 sys.argv) 107 run_tests(args, unittest_argv) 108