1"""Uwb base test."""
2
3import logging
4import re
5
6from mobly import base_test
7from mobly import records
8from mobly import test_runner
9from mobly.controllers import android_device
10
11from test_utils import uwb_test_utils
12
13RELEASE_ID_REGEX = re.compile(r"\w+\.\d+\.\d+")
14
15
16class UwbBaseTest(base_test.BaseTestClass):
17  """Base class for Uwb tests."""
18
19  def setup_class(self):
20    """Sets up the Android devices for Uwb test."""
21    super().setup_class()
22    self.android_devices = self.register_controller(android_device,
23                                                    min_number=2)
24    for ad in self.android_devices:
25      ad.load_snippet("uwb", "com.google.snippet.uwb")
26
27    for ad in self.android_devices:
28      uwb_test_utils.initialize_uwb_country_code_if_not_set(ad)
29
30  def setup_test(self):
31    super().setup_test()
32    for ad in self.android_devices:
33      ad.uwb.logInfo("*** TEST START: " + self.current_test_info.name + " ***")
34
35  def teardown_test(self):
36    super().teardown_test()
37    for ad in self.android_devices:
38      ad.uwb.logInfo("*** TEST END: " + self.current_test_info.name + " ***")
39
40  def teardown_class(self):
41    super().teardown_class()
42    self._record_all()
43
44  def on_fail(self, record):
45    test_name = record.test_name
46    # Single device test
47    if hasattr(self, "dut"):
48      self.dut.take_bug_report(
49          test_name=test_name, destination=self.current_test_info.output_path
50      )
51    else:
52      for count, ad in enumerate(self.android_devices):
53        device_name = "initiator" if not count else "responder"
54        test_device_name = test_name + "_" + device_name
55        ad.take_bug_report(
56            test_name=test_device_name,
57            destination=self.current_test_info.output_path,
58        )
59
60  def _get_effort_name(self) -> str:
61    """Gets the TestTracker effort name from the Android build ID.
62
63    Returns:
64      Testtracker effort name for the test results.
65    """
66    effort_name = "UNKNOWN"
67    for record in self._controller_manager.get_controller_info_records():
68      if record.controller_name == "AndroidDevice" and record.controller_info:
69        build_info = record.controller_info[0]["build_info"]
70        if re.match(RELEASE_ID_REGEX, build_info["build_id"]):
71          effort_name = build_info["build_id"]
72        else:
73          effort_name = build_info[android_device.BuildInfoConstants
74                                   .BUILD_VERSION_INCREMENTAL.build_info_key]
75        break
76    return effort_name
77
78  def _record(self, tr_record: records.TestResultRecord):
79    """Records TestTracker data for a single test.
80
81    Args:
82      tr_record: test case record.
83    """
84    self.record_data({
85        "Test Class": tr_record.test_class,
86        "Test Name": tr_record.test_name,
87        "sponge_properties": {
88            "test_tracker_effort_name": self._get_effort_name(),
89            "test_tracker_uuid": tr_record.uid
90        }
91    })
92
93  def _record_all(self):
94    """Records TestTracker data for all tests."""
95    tr_record_dict = {}
96    for tr_record in self.results.executed:
97      if not tr_record.uid:
98        logging.warning("Missing UID for test %s", tr_record.test_name)
99        continue
100      if tr_record.uid in tr_record_dict:
101        record = tr_record_dict[tr_record.uid]
102        if record in self.results.failed or record in self.results.error:
103          continue
104      tr_record_dict[tr_record.uid] = tr_record
105    for tr_record in tr_record_dict.values():
106      self._record(tr_record)
107
108
109if __name__ == "__main__":
110  test_runner.main()
111