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"""Mobly base test class for Neaby Connections.
16
17Override the NCBaseTestClass#_get_country_code method if the test requires
18a special country code, the 'US' is used by default.
19"""
20
21import logging
22import os
23import time
24
25from mobly import asserts
26from mobly import base_test
27from mobly import records
28from mobly import utils
29from mobly.controllers import android_device
30from mobly.controllers.android_device_lib import errors
31import yaml
32
33from betocq import android_wifi_utils
34from betocq import nc_constants
35from betocq import setup_utils
36from betocq import version
37
38NEARBY_SNIPPET_PACKAGE_NAME = 'com.google.android.nearby.mobly.snippet'
39NEARBY_SNIPPET_2_PACKAGE_NAME = 'com.google.android.nearby.mobly.snippet.second'
40NEARBY_SNIPPET_3P_PACKAGE_NAME = 'com.google.android.nearby.mobly.snippet.thirdparty'
41
42# TODO(b/330803934): Need to design external path for OEM.
43_CONFIG_EXTERNAL_PATH = 'TBD'
44
45
46class NCBaseTestClass(base_test.BaseTestClass):
47  """The Base of Nearby Connection E2E tests."""
48
49  def __init__(self, configs):
50    super().__init__(configs)
51    self.ads: list[android_device.AndroidDevice] = []
52    self.advertiser: android_device.AndroidDevice = None
53    self.discoverer: android_device.AndroidDevice = None
54    self.test_parameters: nc_constants.TestParameters = (
55        nc_constants.TestParameters.from_user_params(self.user_params)
56    )
57    self._test_result_messages: dict[str, str] = {}
58    self._nearby_snippet_apk_path: str = None
59    self._nearby_snippet_2_apk_path: str = None
60    self._nearby_snippet_3p_apk_path: str = None
61    self.performance_test_iterations: int = 1
62    self.num_bug_reports: int = 0
63    self._requires_2_snippet_apks = False
64    self._requires_3p_snippet_apks = False
65    self.__loaded_2_nearby_snippets = False
66    self.__loaded_3p_nearby_snippets = False
67    self.__skipped_test_class = False
68
69  def _get_skipped_test_class_reason(self) -> str | None:
70    return None
71
72  def setup_class(self) -> None:
73    self._setup_openwrt_wifi()
74    self.ads = self.register_controller(android_device, min_number=2)
75    try:
76      self.discoverer = android_device.get_device(
77          self.ads, role='source_device'
78      )
79      self.advertiser = android_device.get_device(
80          self.ads, role='target_device'
81      )
82    except errors.Error:
83      logging.warning(
84          'The source,target devices are not specified in testbed;'
85          'The result may not be expected.'
86      )
87      self.advertiser, self.discoverer = self.ads
88
89    utils.concurrent_exec(
90        self._setup_android_hw_capability,
91        param_list=[[ad] for ad in self.ads],
92        raise_on_exception=True,
93    )
94
95    skipped_test_class_reason = self._get_skipped_test_class_reason()
96    for ad in self.ads:
97      if (
98          not ad.wifi_chipset
99          and self.test_parameters.skip_test_if_wifi_chipset_is_empty
100      ):
101        skipped_test_class_reason = 'wifi_chipset is empty in the config file'
102        ad.log.warning(skipped_test_class_reason)
103    if skipped_test_class_reason:
104      self.__skipped_test_class = True
105      asserts.abort_class(skipped_test_class_reason)
106
107    file_tag = 'files' if 'files' in self.user_params else 'mh_files'
108    self._nearby_snippet_apk_path = self.user_params.get(file_tag, {}).get(
109        'nearby_snippet', ['']
110    )[0]
111    if self.test_parameters.requires_bt_multiplex:
112      self._requires_2_snippet_apks = True
113      self._nearby_snippet_2_apk_path = self.user_params.get(file_tag, {}).get(
114          'nearby_snippet_2', ['']
115      )[0]
116    if self.test_parameters.requires_3p_api_test:
117      self._requires_3p_snippet_apks = True
118      self._nearby_snippet_3p_apk_path = self.user_params.get(file_tag, {}).get(
119          'nearby_snippet_3p', ['']
120      )[0]
121
122    # disconnect from all wifi automatically
123    utils.concurrent_exec(
124        android_wifi_utils.forget_all_wifi,
125        param_list=[[ad] for ad in self.ads],
126        raise_on_exception=True,
127    )
128
129    utils.concurrent_exec(
130        self._setup_android_device,
131        param_list=[[ad] for ad in self.ads],
132        raise_on_exception=True,
133    )
134
135  def _setup_openwrt_wifi(self):
136    """Sets up the wifi connection with OpenWRT."""
137    if not self.user_params.get('use_auto_controlled_wifi_ap', False):
138      return
139
140    self.openwrt = self.register_controller(openwrt_device)[0]
141    if 'wifi_channel' in self.user_params:
142      wifi_channel = self.user_params['wifi_channel']
143      self.wifi_info = self.openwrt.start_wifi(
144          config=wifi_configs.WiFiConfig(
145              channel=wifi_channel,
146              country_code=self._get_country_code(),
147          )
148      )
149    else:
150      wifi_channel = None
151      self.wifi_info = self.openwrt.start_wifi(
152          config=wifi_configs.WiFiConfig(
153              country_code=self._get_country_code(),
154          )
155      )
156
157    if wifi_channel is None:
158      self.test_parameters.wifi_ssid = self.wifi_info.ssid
159      self.test_parameters.wifi_password = self.wifi_info.password
160    elif wifi_channel == 6:
161      self.test_parameters.wifi_2g_ssid = self.wifi_info.ssid
162      self.test_parameters.wifi_2g_password = self.wifi_info.password
163    elif wifi_channel == 36:
164      self.test_parameters.wifi_5g_ssid = self.wifi_info.ssid
165      self.test_parameters.wifi_5g_password = self.wifi_info.password
166    elif wifi_channel == 52:
167      self.test_parameters.wifi_dfs_5g_ssid = self.wifi_info.ssid
168      self.test_parameters.wifi_dfs_5g_password = self.wifi_info.password
169    else:
170      raise ValueError('Unknown Wi-Fi channel: %s' % wifi_channel)
171
172  def _setup_android_hw_capability(
173      self, ad: android_device.AndroidDevice
174  ) -> None:
175    ad.android_version = int(ad.adb.getprop('ro.build.version.release'))
176
177    # TODO(b/330803934): Need to design external path for OEM.
178    if not os.path.isfile(_CONFIG_EXTERNAL_PATH):
179      return
180
181    config_path = _CONFIG_EXTERNAL_PATH
182    with open(config_path, 'r') as f:
183      rule = yaml.safe_load(f).get(ad.model, None)
184      if rule is None:
185        ad.log.warning(f'{ad} Model {ad.model} is not supported in config file')
186        return
187      for key, value in rule.items():
188        ad.log.debug('Setting capability %s to %s', repr(key), repr(value))
189        setattr(ad, key, value)
190
191  def _get_country_code(self) -> str:
192    return 'US'
193
194  def _setup_android_device(self, ad: android_device.AndroidDevice) -> None:
195    ad.debug_tag = ad.serial + '(' + ad.adb.getprop('ro.product.model') + ')'
196    if not ad.is_adb_root:
197      if self.test_parameters.allow_unrooted_device:
198        ad.log.info('Unrooted device is detected. Test coverage is limited')
199      else:
200        asserts.abort_all('The test only can run on rooted device.')
201
202    setup_utils.disable_gms_auto_updates(ad)
203
204    ad.debug_tag = ad.serial + '(' + ad.adb.getprop('ro.product.model') + ')'
205    ad.log.info('try to install nearby_snippet_apk')
206    if self._nearby_snippet_apk_path:
207      setup_utils.install_apk(ad, self._nearby_snippet_apk_path)
208    else:
209      ad.log.warning(
210          'nearby_snippet apk is not specified, '
211          'make sure it is installed in the device'
212      )
213
214    ad.log.info('grant manage external storage permission')
215    setup_utils.grant_manage_external_storage_permission(
216        ad, NEARBY_SNIPPET_PACKAGE_NAME
217    )
218    ad.load_snippet('nearby', NEARBY_SNIPPET_PACKAGE_NAME)
219
220    if self._requires_2_snippet_apks:
221      ad.log.info('try to install nearby_snippet_2_apk')
222      if self._nearby_snippet_2_apk_path:
223        setup_utils.install_apk(ad, self._nearby_snippet_2_apk_path)
224      else:
225        ad.log.warning(
226            'nearby_snippet_2 apk is not specified, '
227            'make sure it is installed in the device'
228        )
229      setup_utils.grant_manage_external_storage_permission(
230          ad, NEARBY_SNIPPET_2_PACKAGE_NAME
231      )
232      ad.load_snippet('nearby2', NEARBY_SNIPPET_2_PACKAGE_NAME)
233      self.__loaded_2_nearby_snippets = True
234    if self._requires_3p_snippet_apks:
235      ad.log.info('try to install nearby_snippet_3p_apk')
236      if self._nearby_snippet_3p_apk_path:
237        setup_utils.install_apk(ad, self._nearby_snippet_3p_apk_path)
238      else:
239        ad.log.warning(
240            'nearby_snippet_3p apk is not specified, '
241            'make sure it is installed in the device'
242        )
243      setup_utils.grant_manage_external_storage_permission(
244          ad, NEARBY_SNIPPET_3P_PACKAGE_NAME
245      )
246      ad.load_snippet('nearby3p', NEARBY_SNIPPET_3P_PACKAGE_NAME)
247      self.__loaded_3p_nearby_snippets = True
248
249    if not ad.nearby.wifiIsEnabled():
250      ad.nearby.wifiEnable()
251    setup_utils.disconnect_from_wifi(ad)
252    setup_utils.enable_logs(ad)
253    setup_utils.disable_redaction(ad)
254    setup_utils.enable_wifi_aware(ad)
255    setup_utils.disable_wlan_deny_list(ad)
256
257    setup_utils.enable_ble_scan_throttling_during_2g_transfer(
258        ad, self.test_parameters.enable_2g_ble_scan_throttling
259    )
260
261    setup_utils.set_country_code(ad, self._get_country_code())
262
263  def setup_test(self):
264    self.record_data({
265        'Test Name': self.current_test_info.name,
266        'properties': {
267            'beto_team': 'Nearby Connections',
268            'beto_feature': 'Nearby Connections',
269        },
270    })
271    self._reset_nearby_connection()
272
273  def _reset_wifi_connection(self) -> None:
274    """Resets wifi connections on both devices."""
275    self.discoverer.nearby.wifiClearConfiguredNetworks()
276    self.advertiser.nearby.wifiClearConfiguredNetworks()
277    time.sleep(nc_constants.WIFI_DISCONNECTION_DELAY.total_seconds())
278
279  def _reset_nearby_connection(self) -> None:
280    """Resets nearby connection."""
281    self.discoverer.nearby.stopDiscovery()
282    self.discoverer.nearby.stopAllEndpoints()
283    self.advertiser.nearby.stopAdvertising()
284    self.advertiser.nearby.stopAllEndpoints()
285    if self.__loaded_2_nearby_snippets:
286      self.discoverer.nearby2.stopDiscovery()
287      self.discoverer.nearby2.stopAllEndpoints()
288      self.advertiser.nearby2.stopAdvertising()
289      self.advertiser.nearby2.stopAllEndpoints()
290    if self.__loaded_3p_nearby_snippets:
291      self.discoverer.nearby3p.stopDiscovery()
292      self.discoverer.nearby3p.stopAllEndpoints()
293      self.advertiser.nearby3p.stopAdvertising()
294      self.advertiser.nearby3p.stopAllEndpoints()
295    time.sleep(nc_constants.NEARBY_RESET_WAIT_TIME.total_seconds())
296
297  def _teardown_device(self, ad: android_device.AndroidDevice) -> None:
298    ad.nearby.transferFilesCleanup()
299    setup_utils.enable_gms_auto_updates(ad)
300
301    if self.test_parameters.disconnect_wifi_after_test:
302      setup_utils.disconnect_from_wifi(ad)
303
304    ad.unload_snippet('nearby')
305    if self.__loaded_2_nearby_snippets:
306      ad.unload_snippet('nearby2')
307    if self.__loaded_3p_nearby_snippets:
308      ad.unload_snippet('nearby3p')
309
310  def teardown_test(self) -> None:
311    utils.concurrent_exec(
312        lambda d: d.services.create_output_excerpts_all(self.current_test_info),
313        param_list=[[ad] for ad in self.ads],
314        raise_on_exception=True,
315    )
316    if hasattr(self, 'openwrt'):
317      self.openwrt.services.create_output_excerpts_all(self.current_test_info)
318
319  def teardown_class(self) -> None:
320    if self.__skipped_test_class:
321      logging.info('Skipping teardown class.')
322      return
323
324    # handle summary results
325    self._summary_test_results()
326
327    utils.concurrent_exec(
328        self._teardown_device,
329        param_list=[[ad] for ad in self.ads],
330        raise_on_exception=True,
331    )
332
333    if hasattr(self, 'openwrt') and hasattr(self, 'wifi_info'):
334      self.openwrt.stop_wifi(self.wifi_info)
335
336  def _summary_test_results(self) -> None:
337    pass
338
339  def on_fail(self, record: records.TestResultRecord) -> None:
340    if self.__skipped_test_class:
341      logging.info('Skipping on_fail.')
342      return
343    if self.test_parameters.skip_bug_report:
344      logging.info('Skipping bug report.')
345      return
346    self.num_bug_reports = self.num_bug_reports + 1
347    if self.num_bug_reports <= nc_constants.MAX_NUM_BUG_REPORT:
348      logging.info('take bug report for failure')
349      android_device.take_bug_reports(
350          self.ads,
351          destination=self.current_test_info.output_path,
352      )
353