1#  Copyright (C) 2024 The Android Open Source Project
2#
3#  Licensed under the Apache License, Version 2.0 (the "License");
4#  you may not use this file except in compliance with the License.
5#  You may obtain a copy of the License at
6#
7#       http://www.apache.org/licenses/LICENSE-2.0
8#
9#  Unless required by applicable law or agreed to in writing, software
10#  distributed under the License is distributed on an "AS IS" BASIS,
11#  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12#  See the License for the specific language governing permissions and
13#  limitations under the License.
14
15"""Base class for BetoCQ test suites."""
16
17import logging
18import os
19
20from mobly import base_suite
21from mobly import records
22import yaml
23
24from betocq import version
25
26
27_BETOCQ_SUITE_NAME = 'BeToCQ'
28
29
30class BaseBetocqSuite(base_suite.BaseSuite):
31  """Base class for BetoCQ test suites.
32
33  Contains methods for aggregating and exporting suite data.
34  """
35
36  def __init__(self, runner, config):
37    super().__init__(runner, config)
38    self._summary_path = None
39    self._summary_writer = None
40
41  def teardown_suite(self):
42    """Collects test class results and reports them as suite properties."""
43    user_data = self._retrieve_user_data_from_summary()
44    class_data = [
45        entry
46        for entry in user_data
47        if records.TestResultEnums.RECORD_CLASS in entry
48        and records.TestResultEnums.RECORD_NAME not in entry
49    ]
50    class_results = {
51      'suite_name': _BETOCQ_SUITE_NAME,
52      'run_identifier': f'v{version.TEST_SCRIPT_VERSION}',
53    }
54    for entry in class_data:
55      properties = entry.get('properties', {})
56      for key, value in properties.items():
57        # prepend '0'/'1' so the properties appear first in lexicographic order
58        if key.endswith('source_device'):
59          if '0_source_device' not in class_results:
60            class_results['0_source_device'] = value
61        if key.endswith('target_device'):
62          if '0_target_device' not in class_results:
63            class_results['0_target_device'] = value
64        if key.endswith('test_result'):
65          class_results[f'1_{entry[records.TestResultEnums.RECORD_CLASS]}'] = (
66              value
67          )
68        if key.endswith('detailed_stats'):
69          class_results[
70              f'1_{entry[records.TestResultEnums.RECORD_CLASS]}_detailed_stats'
71          ] = value
72    self._record_suite_properties(class_results)
73
74  @property
75  def summary_path(self):
76    """Returns the path to the summary file.
77
78    NOTE: This path is only correctly resolved if called within teardown_suite.
79    """
80    if self._summary_path is None:
81      # pylint: disable-next=protected-access
82      self._summary_path = self._runner._test_run_metadata.summary_file_path
83    return self._summary_path
84
85  def _retrieve_user_data_from_summary(self):
86    """Retrieves all user_data entries from the currently streamed summary.
87
88    Use this method to aggregate data written by record_data in test classes.
89
90    NOTE: This method can only be called within teardown_suite.
91
92    Returns:
93      A list of dictionaries, each corresponding to a USER_DATA entry.
94    """
95    if not os.path.isfile(self.summary_path):
96      logging.error(
97          'Cannot retrieve user data for the suite. '
98          'The summary file does not exist: %s',
99          self.summary_path,
100      )
101      return []
102
103    with open(self.summary_path, 'r') as f:
104      return [
105          entry
106          for entry in yaml.safe_load_all(f)
107          if entry['Type'] == records.TestSummaryEntryType.USER_DATA.value
108      ]
109
110  def _record_suite_properties(self, properties):
111    """Record suite properties to the test summary file.
112
113    NOTE: This method can only be called within teardown_suite.
114
115    Args:
116      properties: dict, the properties to add to the summary
117    """
118    if self._summary_writer is None:
119      self._summary_writer = records.TestSummaryWriter(self.summary_path)
120    content = {'properties': properties}
121    self._summary_writer.dump(content, records.TestSummaryEntryType.USER_DATA)
122