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