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# Lint as: python3
16"""CTS-V Wi-Fi Aware test reimplemented in Mobly."""
17import datetime
18import enum
19import logging
20import random
21import sys
22import time
23from typing import Tuple, Any
24
25from mobly import asserts
26from mobly import base_test
27from mobly import records
28from mobly import test_runner
29from mobly import utils
30from mobly.controllers import android_device
31from mobly.controllers.android_device_lib import callback_handler_v2
32from mobly.snippet import callback_event
33import wifi_test_utils
34
35from aware import constants
36from aware import aware_lib_utils
37
38PACKAGE_NAME = constants.WIFI_SNIPPET_PACKAGE_NAME
39_DEFAULT_TIMEOUT = constants.WAIT_WIFI_STATE_TIME_OUT.total_seconds()
40_REQUEST_NETWORK_TIMEOUT_MS = 15 * 1000
41_MSG_ID_SUB_TO_PUB = random.randint(1000, 5000)
42_MSG_ID_PUB_TO_SUB = random.randint(5001, 9999)
43_MSG_SUB_TO_PUB = "Let's talk [Random Identifier: %s]" % utils.rand_ascii_str(5)
44_MSG_PUB_TO_SUB = 'Ready [Random Identifier: %s]' % utils.rand_ascii_str(5)
45_PUB_SSI = constants.WifiAwareTestConstants.PUB_SSI
46_MATCH_FILTER = [constants.WifiAwareTestConstants.MATCH_FILTER_BYTES]
47_CALLBACK_NAME = constants.DiscoverySessionCallbackParamsType.CALLBACK_NAME
48_IS_SESSION_INIT = constants.DiscoverySessionCallbackParamsType.IS_SESSION_INIT
49_TRANSPORT_TYPE_WIFI_AWARE = (
50    constants.NetworkCapabilities.Transport.TRANSPORT_WIFI_AWARE
51)
52_LARGE_ENOUGH_DISTANCE_MM = 100000  # 100 meters
53_MIN_RSSI = -100
54_WAIT_SEC_FOR_RTT_INITIATOR_RESPONDER_SWITCH = 5
55
56
57@enum.unique
58class AttachCallBackMethodType(enum.StrEnum):
59    """Represents Attach Callback Method Type in Wi-Fi Aware.
60
61    https://developer.android.com/reference/android/net/wifi/aware/AttachCallback
62    """
63    ATTACHED = 'onAttached'
64    ATTACH_FAILED = 'onAttachFailed'
65    AWARE_SESSION_TERMINATED = 'onAwareSessionTerminated'
66
67
68class WifiAwareManagerTest(base_test.BaseTestClass):
69    """Wi-Fi Aware test class."""
70
71    ads: list[android_device.AndroidDevice]
72    publisher: android_device.AndroidDevice
73    subscriber: android_device.AndroidDevice
74
75    # Wi-Fi Aware attach session ID
76    publisher_attach_session: str | None = None
77    subscriber_attach_session: str | None = None
78    # Wi-Fi Aware discovery session ID
79    publish_session: str | None = None
80    subscribe_session: str | None = None
81    # Wi-Fi Aware peer ID
82    publisher_peer: int | None = None
83    subscriber_peer: int | None = None
84    # Mac addresses.
85    publisher_mac: str | None = None
86    subscriber_mac: str | None = None
87
88    def setup_class(self):
89        # Register and set up Android devices in parallel.
90        self.ads = self.register_controller(android_device, min_number=2)
91        self.publisher = self.ads[0]
92        self.subscriber = self.ads[1]
93
94        def setup_device(device: android_device.AndroidDevice):
95            device.load_snippet(
96                'wifi_aware_snippet', PACKAGE_NAME
97            )
98            aware_lib_utils.control_wifi(device, wifi_state=True)
99            asserts.abort_all_if(
100                not device.wifi_aware_snippet.wifiAwareIsSupported(),
101                f'{device} does not support Wi-Fi Aware.',
102            )
103            wifi_test_utils.set_screen_on_and_unlock(device)
104            asserts.abort_all_if(
105                not device.wifi_aware_snippet.wifiAwareIsAvailable(),
106                f'{device} Wi-Fi Aware is not available.',
107            )
108
109        utils.concurrent_exec(
110            setup_device,
111            ((self.publisher,), (self.subscriber,)),
112            max_workers=2,
113            raise_on_exception=True,
114        )
115
116    def test_data_path_open_unsolicited_pub_and_passive_sub(self) -> None:
117        """Test OPEN Wi-Fi Aware network with unsolicited publish and passive subscribe.
118
119        Steps:
120        1. Attach a Wi-Fi Aware session on each device.
121        2. Publish and subscribe to a discovery session.
122        3. Send messages through discovery session’s API.
123        4. Request a Wi-Fi Aware network.
124        5. Establish a socket connection and send messages through it.
125        """
126
127        self._test_wifi_aware(
128            pub_config=constants.PublishConfig(
129                service_specific_info=constants.WifiAwareTestConstants.PUB_SSI,
130                match_filter=[constants.WifiAwareTestConstants.MATCH_FILTER_BYTES],
131                publish_type=constants.PublishType.UNSOLICITED,
132                ranging_enabled=False,
133            ),
134            sub_config=constants.SubscribeConfig(
135                service_specific_info=constants.WifiAwareTestConstants.SUB_SSI,
136                match_filter=[constants.WifiAwareTestConstants.MATCH_FILTER_BYTES],
137                subscribe_type=constants.SubscribeType.PASSIVE,
138            ),
139        )
140
141    def test_data_path_passphrase_unsolicited_pub_and_passive_sub(self) -> None:
142        """Test Wi-Fi Aware network with passphrase, unsolicited publish, and passive subscribe.
143
144        Steps:
145        1. Attach a Wi-Fi Aware session on each device.
146        2. Publish and subscribe to a discovery session.
147        3. Send messages through discovery session’s API.
148        4. Request a Wi-Fi Aware network.
149        5. Establish a socket connection and send messages through it.
150        """
151
152        self._test_wifi_aware(
153            pub_config=constants.PublishConfig(
154                service_specific_info=constants.WifiAwareTestConstants.PUB_SSI,
155                match_filter=[constants.WifiAwareTestConstants.MATCH_FILTER_BYTES],
156                publish_type=constants.PublishType.UNSOLICITED,
157                ranging_enabled=False,
158            ),
159            sub_config=constants.SubscribeConfig(
160                service_specific_info=constants.WifiAwareTestConstants.SUB_SSI,
161                match_filter=[constants.WifiAwareTestConstants.MATCH_FILTER_BYTES],
162                subscribe_type=constants.SubscribeType.PASSIVE,
163            ),
164            network_specifier_on_pub=constants.WifiAwareNetworkSpecifier(
165                psk_passphrase=constants.WifiAwareTestConstants.PASSWORD,
166                transport_protocol=constants.WifiAwareTestConstants.TRANSPORT_PROTOCOL_TCP,
167            ),
168            network_specifier_on_sub=constants.WifiAwareNetworkSpecifier(
169                psk_passphrase=constants.WifiAwareTestConstants.PASSWORD,
170            )
171        )
172
173    def test_data_path_pmk_unsolicited_pub_and_passive_sub(self) -> None:
174        """Test Wi-Fi Aware network using PMK with unsolicited publish and passive subscribe.
175
176        Steps:
177        1. Attach a Wi-Fi Aware session on each device.
178        2. Publish and subscribe to a discovery session.
179        3. Send messages through discovery session’s API.
180        4. Request a Wi-Fi Aware network.
181        5. Establish a socket connection and send messages through it.
182        """
183
184        self._test_wifi_aware(
185            pub_config=constants.PublishConfig(
186                service_specific_info=constants.WifiAwareTestConstants.PUB_SSI,
187                match_filter=[constants.WifiAwareTestConstants.MATCH_FILTER_BYTES],
188                publish_type=constants.PublishType.UNSOLICITED,
189                ranging_enabled=False,
190            ),
191            sub_config=constants.SubscribeConfig(
192                service_specific_info=constants.WifiAwareTestConstants.SUB_SSI,
193                match_filter=[constants.WifiAwareTestConstants.MATCH_FILTER_BYTES],
194                subscribe_type=constants.SubscribeType.PASSIVE,
195            ),
196            network_specifier_on_pub=constants.WifiAwareNetworkSpecifier(
197                transport_protocol=constants.WifiAwareTestConstants.TRANSPORT_PROTOCOL_TCP,
198                pmk=constants.WifiAwareTestConstants.PMK,
199            ),
200            network_specifier_on_sub=constants.WifiAwareNetworkSpecifier(
201                data_path_security_config=constants.WifiAwareDataPathSecurityConfig(
202                    pmk=constants.WifiAwareTestConstants.PMK
203                )
204            )
205        )
206
207    def test_data_path_open_solicited_pub_and_active_sub(self) -> None:
208        """Test OPEN Wi-Fi Aware network with solicited publish and active subscribe.
209
210        Steps:
211        1. Attach a Wi-Fi Aware session on each device.
212        2. Publish and subscribe to a discovery session.
213        3. Send messages through discovery session’s API.
214        4. Request a Wi-Fi Aware network.
215        5. Establish a socket connection and send messages through it.
216        """
217
218        self._test_wifi_aware(
219            pub_config=constants.PublishConfig(
220                service_specific_info=constants.WifiAwareTestConstants.PUB_SSI,
221                match_filter=[constants.WifiAwareTestConstants.MATCH_FILTER_BYTES],
222                publish_type=constants.PublishType.SOLICITED,
223                ranging_enabled=False,
224            ),
225            sub_config=constants.SubscribeConfig(
226                service_specific_info=constants.WifiAwareTestConstants.SUB_SSI,
227                match_filter=[constants.WifiAwareTestConstants.MATCH_FILTER_BYTES],
228                subscribe_type=constants.SubscribeType.ACTIVE,
229            ),
230
231        )
232
233    def test_data_path_passphrase_solicited_pub_and_active_sub(self) -> None:
234        """Test password-protected Wi-Fi Aware network with solicited publish and active subscribe.
235
236        Steps:
237        1. Attach a Wi-Fi Aware session on each device.
238        2. Publish and subscribe to a discovery session.
239        3. Send messages through discovery session’s API.
240        4. Request a Wi-Fi Aware network.
241        5. Establish a socket connection and send messages through it.
242        """
243
244        self._test_wifi_aware(
245            pub_config=constants.PublishConfig(
246                service_specific_info=constants.WifiAwareTestConstants.PUB_SSI,
247                match_filter=[constants.WifiAwareTestConstants.MATCH_FILTER_BYTES],
248                publish_type=constants.PublishType.SOLICITED,
249                ranging_enabled=False,
250            ),
251            sub_config=constants.SubscribeConfig(
252                service_specific_info=constants.WifiAwareTestConstants.SUB_SSI,
253                match_filter=[constants.WifiAwareTestConstants.MATCH_FILTER_BYTES],
254                subscribe_type=constants.SubscribeType.ACTIVE,
255            ),
256            network_specifier_on_pub=constants.WifiAwareNetworkSpecifier(
257                psk_passphrase=constants.WifiAwareTestConstants.PASSWORD,
258                transport_protocol=constants.WifiAwareTestConstants.TRANSPORT_PROTOCOL_TCP,
259            ),
260            network_specifier_on_sub=constants.WifiAwareNetworkSpecifier(
261                psk_passphrase=constants.WifiAwareTestConstants.PASSWORD,
262            )
263        )
264
265    def test_data_path_pmk_solicited_pub_and_active_sub(self) -> None:
266        """Test Wi-Fi Aware network using PMK with solicited publish and active subscribe.
267
268        Steps:
269        1. Attach a Wi-Fi Aware session on each device.
270        2. Publish and subscribe to a discovery session.
271        3. Send messages through discovery session’s API.
272        4. Request a Wi-Fi Aware network.
273        5. Establish a socket connection and send messages through it.
274        """
275
276        self._test_wifi_aware(
277            pub_config=constants.PublishConfig(
278                service_specific_info=constants.WifiAwareTestConstants.PUB_SSI,
279                match_filter=[constants.WifiAwareTestConstants.MATCH_FILTER_BYTES],
280                publish_type=constants.PublishType.SOLICITED,
281                ranging_enabled=False,
282            ),
283            sub_config=constants.SubscribeConfig(
284                service_specific_info=constants.WifiAwareTestConstants.SUB_SSI,
285                match_filter=[constants.WifiAwareTestConstants.MATCH_FILTER_BYTES],
286                subscribe_type=constants.SubscribeType.ACTIVE,
287            ),
288            network_specifier_on_pub=constants.WifiAwareNetworkSpecifier(
289                transport_protocol=constants.WifiAwareTestConstants.TRANSPORT_PROTOCOL_TCP,
290                pmk=constants.WifiAwareTestConstants.PMK,
291            ),
292            network_specifier_on_sub=constants.WifiAwareNetworkSpecifier(
293                data_path_security_config=constants.WifiAwareDataPathSecurityConfig(
294                    pmk=constants.WifiAwareTestConstants.PMK
295                )
296            )
297        )
298
299    def test_data_path_open_unsolicited_pub_accept_any_and_passive_sub(self) -> None:
300        """Test OPEN Wi-Fi Aware with unsolicited publish (accept any peer) and passive subscribe.
301
302        Steps:
303        1. Attach a Wi-Fi Aware session on each device.
304        2. Publish and subscribe to a discovery session.
305        3. Send messages through discovery session’s API.
306        4. Request a Wi-Fi Aware network.
307        5. Establish a socket connection and send messages through it.
308        """
309        self._test_wifi_aware(
310            pub_config=constants.PublishConfig(
311                service_specific_info=constants.WifiAwareTestConstants.PUB_SSI,
312                match_filter=[constants.WifiAwareTestConstants.MATCH_FILTER_BYTES],
313                publish_type=constants.PublishType.UNSOLICITED,
314                ranging_enabled=False,
315            ),
316            sub_config=constants.SubscribeConfig(
317                service_specific_info=constants.WifiAwareTestConstants.SUB_SSI,
318                match_filter=[constants.WifiAwareTestConstants.MATCH_FILTER_BYTES],
319                subscribe_type=constants.SubscribeType.PASSIVE,
320            ),
321            is_pub_accept_any_peer=True,
322        )
323
324    def test_data_path_passphrase_unsolicited_pub_accept_any_and_passive_sub(self) -> None:
325        """Test Wi-Fi Aware with passphrase unsolicited publish (accept any), and passive subscribe.
326
327        Steps:
328        1. Attach a Wi-Fi Aware session on each device.
329        2. Publish and subscribe to a discovery session.
330        3. Send messages through discovery session’s API.
331        4. Request a Wi-Fi Aware network.
332        5. Establish a socket connection and send messages through it.
333        """
334        self._test_wifi_aware(
335            pub_config=constants.PublishConfig(
336                service_specific_info=constants.WifiAwareTestConstants.PUB_SSI,
337                match_filter=[constants.WifiAwareTestConstants.MATCH_FILTER_BYTES],
338                publish_type=constants.PublishType.UNSOLICITED,
339                ranging_enabled=False,
340            ),
341            sub_config=constants.SubscribeConfig(
342                service_specific_info=constants.WifiAwareTestConstants.SUB_SSI,
343                match_filter=[constants.WifiAwareTestConstants.MATCH_FILTER_BYTES],
344                subscribe_type=constants.SubscribeType.PASSIVE,
345            ),
346            network_specifier_on_pub=constants.WifiAwareNetworkSpecifier(
347                psk_passphrase=constants.WifiAwareTestConstants.PASSWORD,
348                transport_protocol=constants.WifiAwareTestConstants.TRANSPORT_PROTOCOL_TCP,
349            ),
350            network_specifier_on_sub=constants.WifiAwareNetworkSpecifier(
351                psk_passphrase=constants.WifiAwareTestConstants.PASSWORD,
352            ),
353            is_pub_accept_any_peer=True,
354        )
355
356    def test_data_path_pmk_unsolicited_pub_accept_any_and_passive_sub(self) -> None:
357        """Test Wi-Fi Aware with PMK, unsolicited publish (accept any), and passive subscribe.
358
359        Steps:
360        1. Attach a Wi-Fi Aware session on each device.
361        2. Publish and subscribe to a discovery session.
362        3. Send messages through discovery session’s API.
363        4. Request a Wi-Fi Aware network.
364        5. Establish a socket connection and send messages through it.
365        """
366        self._test_wifi_aware(
367            pub_config=constants.PublishConfig(
368                service_specific_info=constants.WifiAwareTestConstants.PUB_SSI,
369                match_filter=[constants.WifiAwareTestConstants.MATCH_FILTER_BYTES],
370                publish_type=constants.PublishType.UNSOLICITED,
371                ranging_enabled=False,
372            ),
373            sub_config=constants.SubscribeConfig(
374                service_specific_info=constants.WifiAwareTestConstants.SUB_SSI,
375                match_filter=[constants.WifiAwareTestConstants.MATCH_FILTER_BYTES],
376                subscribe_type=constants.SubscribeType.PASSIVE,
377            ),
378            network_specifier_on_pub=constants.WifiAwareNetworkSpecifier(
379                transport_protocol=constants.WifiAwareTestConstants.TRANSPORT_PROTOCOL_TCP,
380                pmk=constants.WifiAwareTestConstants.PMK,
381            ),
382            network_specifier_on_sub=constants.WifiAwareNetworkSpecifier(
383                data_path_security_config=constants.WifiAwareDataPathSecurityConfig(
384                    pmk=constants.WifiAwareTestConstants.PMK
385                )
386            ),
387            is_pub_accept_any_peer=True,
388        )
389
390    def test_data_path_open_solicited_pub_accept_any_active_sub(self) -> None:
391        """Test Wi-Fi Aware with open network, solicited publish (accept any), and active subscribe.
392
393        Steps:
394        1. Attach a Wi-Fi Aware session on each device.
395        2. Publish and subscribe to a discovery session.
396        3. Send messages through discovery session’s API.
397        4. Request a Wi-Fi Aware network.
398        5. Establish a socket connection and send messages through it.
399        """
400        self._test_wifi_aware(
401            pub_config=constants.PublishConfig(
402                service_specific_info=constants.WifiAwareTestConstants.PUB_SSI,
403                match_filter=[constants.WifiAwareTestConstants.MATCH_FILTER_BYTES],
404                publish_type=constants.PublishType.SOLICITED,
405                ranging_enabled=False,
406            ),
407            sub_config=constants.SubscribeConfig(
408                service_specific_info=constants.WifiAwareTestConstants.SUB_SSI,
409                match_filter=[constants.WifiAwareTestConstants.MATCH_FILTER_BYTES],
410                subscribe_type=constants.SubscribeType.ACTIVE,
411            ),
412            is_pub_accept_any_peer=True,
413        )
414
415    def test_data_passphrase_solicited_pub_accept_any_and_active_sub(self) -> None:
416        """Test Wi-Fi Aware with passphrase, solicited publish (accept any), and active subscribe.
417
418        Steps:
419        1. Attach a Wi-Fi Aware session on each device.
420        2. Publish and subscribe to a discovery session.
421        3. Send messages through discovery session’s API.
422        4. Request a Wi-Fi Aware network.
423        5. Establish a socket connection and send messages through it.
424        """
425        self._test_wifi_aware(
426            pub_config=constants.PublishConfig(
427                service_specific_info=constants.WifiAwareTestConstants.PUB_SSI,
428                match_filter=[constants.WifiAwareTestConstants.MATCH_FILTER_BYTES],
429                publish_type=constants.PublishType.SOLICITED,
430                ranging_enabled=False,
431            ),
432            sub_config=constants.SubscribeConfig(
433                service_specific_info=constants.WifiAwareTestConstants.SUB_SSI,
434                match_filter=[constants.WifiAwareTestConstants.MATCH_FILTER_BYTES],
435                subscribe_type=constants.SubscribeType.ACTIVE,
436            ),
437            network_specifier_on_pub=constants.WifiAwareNetworkSpecifier(
438                psk_passphrase=constants.WifiAwareTestConstants.PASSWORD,
439                transport_protocol=constants.WifiAwareTestConstants.TRANSPORT_PROTOCOL_TCP,
440            ),
441            network_specifier_on_sub=constants.WifiAwareNetworkSpecifier(
442                psk_passphrase=constants.WifiAwareTestConstants.PASSWORD,
443            ),
444            is_pub_accept_any_peer=True,
445        )
446
447    def test_data_path_pmk_solicited_pub_accept_any_and_active_sub(self) -> None:
448        """Test Wi-Fi Aware with PMK, solicited publish (accept any), and active subscribe.
449
450        Steps:
451        1. Attach a Wi-Fi Aware session on each device.
452        2. Publish and subscribe to a discovery session.
453        3. Send messages through discovery session’s API.
454        4. Request a Wi-Fi Aware network.
455        5. Establish a socket connection and send messages through it.
456        """
457
458        self._test_wifi_aware(
459            pub_config=constants.PublishConfig(
460                service_specific_info=constants.WifiAwareTestConstants.PUB_SSI,
461                match_filter=[constants.WifiAwareTestConstants.MATCH_FILTER_BYTES],
462                publish_type=constants.PublishType.SOLICITED,
463                ranging_enabled=False,
464            ),
465            sub_config=constants.SubscribeConfig(
466                service_specific_info=constants.WifiAwareTestConstants.SUB_SSI,
467                match_filter=[constants.WifiAwareTestConstants.MATCH_FILTER_BYTES],
468                subscribe_type=constants.SubscribeType.ACTIVE,
469            ),
470            network_specifier_on_pub=constants.WifiAwareNetworkSpecifier(
471                transport_protocol=constants.WifiAwareTestConstants.TRANSPORT_PROTOCOL_TCP,
472                pmk=constants.WifiAwareTestConstants.PMK,
473            ),
474            network_specifier_on_sub=constants.WifiAwareNetworkSpecifier(
475                data_path_security_config=constants.WifiAwareDataPathSecurityConfig(
476                    pmk=constants.WifiAwareTestConstants.PMK
477                )
478            ),
479            is_pub_accept_any_peer=True,
480        )
481
482    def test_discovery_ranging_to_peer_handle(self) -> None:
483        """Test ranging to a Wi-Fi Aware peer handle.
484
485        Steps:
486        1. Attach a Wi-Fi Aware session on each device.
487        2. Publish and subscribe to a discovery session.
488        3. Send messages through discovery session’s API.
489        4. Test ranging to Wi-Fi Aware peer handle.
490        """
491        # Check test condition.
492        self._skip_if_wifi_rtt_is_not_supported()
493
494        # Step 1 - 3. Publish and subscribe Wi-Fi Aware service.
495        self._publish_and_subscribe(
496            pub_config=constants.PublishConfig(
497                publish_type=constants.PublishType.UNSOLICITED,
498                ranging_enabled=True,
499            ),
500            sub_config=constants.SubscribeConfig(
501                subscribe_type=constants.SubscribeType.PASSIVE,
502                max_distance_mm=_LARGE_ENOUGH_DISTANCE_MM,
503            ),
504        )
505
506        # 4. Perform ranging on the publisher and subscriber, respectively.
507        self.publisher.log.info(
508            'Performing ranging to peer ID %d.',  self.publisher_peer
509        )
510        self._perform_ranging(
511            self.publisher,
512            constants.RangingRequest(peer_ids=[self.publisher_peer]),
513        )
514
515        # RTT initiator/responder role switch takes time. We don't have an
516        # API to enforce it. So wait a few seconds for a semi-arbitrary
517        # teardown.
518        time.sleep(_WAIT_SEC_FOR_RTT_INITIATOR_RESPONDER_SWITCH)
519        self.subscriber.log.info(
520            'Performing ranging to peer ID %d.', self.subscriber_peer
521        )
522        self._perform_ranging(
523            self.subscriber,
524            constants.RangingRequest(peer_ids=[self.subscriber_peer]),
525        )
526
527    def test_discovery_ranging_to_peer_mac_address(self) -> None:
528        """Test ranging to a Wi-Fi Aware peer mac address.
529
530        Steps:
531        1. Attach a Wi-Fi Aware session on each device.
532        2. Publish and subscribe to a discovery session.
533        3. Send messages through discovery session’s API.
534        4. Test ranging to Wi-Fi Aware peer mac address.
535        """
536        # Check test condition.
537        self._skip_if_wifi_rtt_is_not_supported()
538
539        # Step 1 - 3. Publish and subscribe Wi-Fi Aware service.
540        self._publish_and_subscribe(
541            pub_config=constants.PublishConfig(
542                publish_type=constants.PublishType.UNSOLICITED,
543                ranging_enabled=True,
544            ),
545            sub_config=constants.SubscribeConfig(
546                subscribe_type=constants.SubscribeType.PASSIVE,
547                max_distance_mm=_LARGE_ENOUGH_DISTANCE_MM,
548            ),
549        )
550
551        # 4. Perform ranging on the publisher and subscriber, respectively.
552        self.publisher.log.info(
553            'Performing ranging to peer MAC address %s.', self.subscriber_mac
554        )
555        self._perform_ranging(
556            self.publisher,
557            constants.RangingRequest(peer_mac_addresses=[self.subscriber_mac]),
558        )
559
560        # RTT initiator/responder role switch takes time. We don't have an
561        # API to enforce it. So wait a few seconds for a semi-arbitrary
562        # teardown.
563        time.sleep(_WAIT_SEC_FOR_RTT_INITIATOR_RESPONDER_SWITCH)
564        self.subscriber.log.info(
565            'Performing ranging to peer MAC address %s.', self.publisher_mac
566        )
567        self._perform_ranging(
568            self.subscriber,
569            constants.RangingRequest(peer_mac_addresses=[self.publisher_mac]),
570        )
571
572    def test_data_path_force_channel_setup(self):
573        """ Test Wi-Fi Aware with PMK, force channel publish, and subscribe.
574
575        Steps:
576        1. Attach a Wi-Fi Aware session on each device.
577        2. Publish and subscribe to a discovery session.
578        3. Send messages through discovery session’s API.
579        4. Request a Wi-Fi Aware network.
580        5. Establish a socket connection and send messages through it.
581        """
582        # The support of this function depends on the chip used.
583        asserts.skip_if(
584            not self.publisher.wifi_aware_snippet.wifiAwareIsSetChannelOnDataPathSupported(),
585            'Publish device not support this test feature.'
586        )
587        asserts.skip_if(
588            not self.subscriber.wifi_aware_snippet.wifiAwareIsSetChannelOnDataPathSupported(),
589            'Subscriber device not support this test feature.'
590        )
591
592        self._test_wifi_aware(
593            pub_config=constants.PublishConfig(
594                service_specific_info=constants.WifiAwareTestConstants.PUB_SSI,
595                match_filter=[constants.WifiAwareTestConstants.MATCH_FILTER_BYTES],
596                publish_type=constants.PublishType.UNSOLICITED,
597                ranging_enabled=False,
598            ),
599            sub_config=constants.SubscribeConfig(
600                service_specific_info=constants.WifiAwareTestConstants.SUB_SSI,
601                match_filter=[constants.WifiAwareTestConstants.MATCH_FILTER_BYTES],
602                subscribe_type=constants.SubscribeType.PASSIVE,
603            ),
604            network_specifier_on_pub=constants.WifiAwareNetworkSpecifier(
605                transport_protocol=constants.WifiAwareTestConstants.TRANSPORT_PROTOCOL_TCP,
606                pmk=constants.WifiAwareTestConstants.PMK,
607                channel_frequency_m_hz=constants.WifiAwareTestConstants.CHANNEL_IN_MHZ,
608            ),
609            network_specifier_on_sub=constants.WifiAwareNetworkSpecifier(
610                data_path_security_config=constants.WifiAwareDataPathSecurityConfig(
611                    pmk=constants.WifiAwareTestConstants.PMK
612                ),
613                channel_frequency_m_hz=constants.WifiAwareTestConstants.CHANNEL_IN_MHZ,
614            )
615        )
616
617    def _test_wifi_aware(
618        self,
619        pub_config: constants.PublishConfig,
620        sub_config: constants.SubscribeConfig,
621        network_specifier_on_pub: constants.WifiAwareNetworkSpecifier | None = None,
622        network_specifier_on_sub: constants.WifiAwareNetworkSpecifier | None = None,
623        is_pub_accept_any_peer: bool = False,
624    ):
625        """Tests Wi-Fi Aware using given publish and subscribe configurations."""
626        # Step 1 - 3: Publish and subscribe Wi-Fi Aware service and send
627        # messages through a Wi-Fi Aware session.
628        self._publish_and_subscribe(pub_config, sub_config)
629
630        # 4. Request a Wi-Fi Aware network.
631        pub_accept_handler = self.publisher.wifi_aware_snippet.connectivityServerSocketAccept()
632        network_id = pub_accept_handler.callback_id
633        pub_local_port = pub_accept_handler.ret_value
634        if network_specifier_on_pub and (
635            network_specifier_on_pub.psk_passphrase or network_specifier_on_pub.pmk):
636            network_specifier_on_pub.port = pub_local_port
637        pub_network_cb_handler = self._request_network(
638            ad=self.publisher,
639            discovery_session=self.publish_session,
640            peer=self.publisher_peer,
641            net_work_request_id=network_id,
642            network_specifier_params=network_specifier_on_pub,
643            is_accept_any_peer=is_pub_accept_any_peer,
644        )
645        sub_network_cb_handler = self._request_network(
646            ad=self.subscriber,
647            discovery_session=self.subscribe_session,
648            peer=self.subscriber_peer,
649            net_work_request_id=network_id,
650            network_specifier_params=network_specifier_on_sub,
651        )
652        expected_channel = None
653        if (network_specifier_on_pub and
654                network_specifier_on_pub.channel_frequency_m_hz and
655                network_specifier_on_sub and
656                network_specifier_on_sub.channel_frequency_m_hz):
657            expected_channel = network_specifier_on_sub.channel_frequency_m_hz
658        self._wait_for_network(
659            ad=self.publisher,
660            request_network_cb_handler=pub_network_cb_handler,
661            expected_channel=None,
662        )
663        self._wait_for_network(
664            ad=self.subscriber,
665            request_network_cb_handler=sub_network_cb_handler,
666            expected_channel=expected_channel,
667        )
668        # 5. Establish a socket connection and send messages through it.
669        self._establish_socket_and_send_msg(
670            pub_accept_handler=pub_accept_handler,
671            network_id=network_id,
672            pub_local_port=pub_local_port
673
674        )
675        self.publisher.wifi_aware_snippet.connectivityUnregisterNetwork(
676            network_id
677        )
678        self.subscriber.wifi_aware_snippet.connectivityUnregisterNetwork(
679            network_id
680        )
681        self.publisher.wifi_aware_snippet.wifiAwareCloseDiscoverSession(
682            self.publish_session
683        )
684        self.publish_session = None
685        self.subscriber.wifi_aware_snippet.wifiAwareCloseDiscoverSession(
686            self.subscribe_session
687        )
688        self.subscribe_session = None
689        self.publisher.wifi_aware_snippet.wifiAwareDetach(
690            self.publisher_attach_session
691        )
692        self.publisher_attach_session = None
693        self.subscriber.wifi_aware_snippet.wifiAwareDetach(
694            self.subscriber_attach_session
695        )
696        self.subscriber_attach_session = None
697        self.publisher.wifi_aware_snippet.connectivityCloseAllSocket(network_id)
698        self.subscriber.wifi_aware_snippet.connectivityCloseAllSocket(network_id)
699
700    def _publish_and_subscribe(self, pub_config, sub_config):
701        """Publishes and subscribes a Wi-Fi Aware session."""
702        # 1. Attach Wi-Fi Aware sessions.
703        self.publisher_attach_session, self.publisher_mac = (
704            self._start_attach(
705                self.publisher, is_ranging_enabled=pub_config.ranging_enabled
706            )
707        )
708        self.subscriber_attach_session, self.subscriber_mac = (
709            self._start_attach(
710                self.subscriber, is_ranging_enabled=pub_config.ranging_enabled
711            )
712        )
713
714        # 2.1. Initialize discovery sessions (publish and subscribe).
715        pub_aware_session_cb_handler = self._start_publish(
716            attach_session_id=self.publisher_attach_session,
717            pub_config=pub_config,
718        )
719        self.publish_session = pub_aware_session_cb_handler.callback_id
720        self.publisher.log.info('Created the publish session.')
721        sub_aware_session_cb_handler = self._start_subscribe(
722            attach_session_id=self.subscriber_attach_session,
723            sub_config=sub_config,
724        )
725        self.subscribe_session = sub_aware_session_cb_handler.callback_id
726        self.subscriber.log.info('Subscribe session created.')
727        # 2.2. Wait for discovery.
728        self.subscriber_peer = self._wait_for_discovery(
729            sub_aware_session_cb_handler,
730            pub_service_specific_info=pub_config.service_specific_info,
731            is_ranging_enabled=pub_config.ranging_enabled,
732        )
733        self.subscriber.log.info('Subscriber discovered the published service.')
734        # 3. Send messages through discovery session’s API.
735        self.publisher_peer = self._send_msg_through_discovery_session(
736            sender=self.subscriber,
737            sender_aware_session_cb_handler=sub_aware_session_cb_handler,
738            receiver=self.publisher,
739            receiver_aware_session_cb_handler=pub_aware_session_cb_handler,
740            discovery_session=self.subscribe_session,
741            peer=self.subscriber_peer,
742            send_message=_MSG_SUB_TO_PUB,
743            send_message_id=_MSG_ID_SUB_TO_PUB,
744        )
745        logging.info(
746            'The subscriber sent a message and the publisher received it.'
747        )
748        self._send_msg_through_discovery_session(
749            sender=self.publisher,
750            sender_aware_session_cb_handler=pub_aware_session_cb_handler,
751            receiver=self.subscriber,
752            receiver_aware_session_cb_handler=sub_aware_session_cb_handler,
753            discovery_session=self.publish_session,
754            peer=self.publisher_peer,
755            send_message=_MSG_PUB_TO_SUB,
756            send_message_id=_MSG_ID_PUB_TO_SUB,
757        )
758        logging.info(
759            'The publisher sent a message and the subscriber received it.'
760        )
761
762    def _establish_socket_and_send_msg(
763        self,
764        pub_accept_handler: callback_handler_v2.CallbackHandlerV2,
765        network_id: str,
766        pub_local_port: int
767    ):
768        """Handles socket-based communication between publisher and subscriber."""
769        # Init socket
770        # Create a ServerSocket and makes it listen for client connections.
771        self.subscriber.wifi_aware_snippet.connectivityCreateSocketOverWiFiAware(
772            network_id, pub_local_port
773        )
774        self._wait_accept_success(pub_accept_handler)
775        # Subscriber Send socket data
776        self.subscriber.log.info('Subscriber create a socket.')
777        self._send_socket_msg(
778            sender_ad=self.subscriber,
779            receiver_ad=self.publisher,
780            msg=constants.WifiAwareTestConstants.MSG_CLIENT_TO_SERVER,
781            send_callback_id=network_id,
782            receiver_callback_id=network_id
783        )
784        self._send_socket_msg(
785            sender_ad=self.publisher,
786            receiver_ad=self.subscriber,
787            msg=constants.WifiAwareTestConstants.MSG_SERVER_TO_CLIENT,
788            send_callback_id=network_id,
789            receiver_callback_id=network_id
790        )
791        self.publisher.wifi_aware_snippet.connectivityCloseWrite(network_id)
792        self.subscriber.wifi_aware_snippet.connectivityCloseWrite(network_id)
793        self.publisher.wifi_aware_snippet.connectivityCloseRead(network_id)
794        self.subscriber.wifi_aware_snippet.connectivityCloseRead(network_id)
795        logging.info('Communicated through socket connection of Wi-Fi Aware network successfully.')
796
797    def _wait_accept_success(
798        self,
799        pub_accept_handler: callback_handler_v2.CallbackHandlerV2
800    ) -> None:
801        pub_accept_event = pub_accept_handler.waitAndGet(
802            event_name=constants.SnippetEventNames.SERVER_SOCKET_ACCEPT,
803            timeout=_DEFAULT_TIMEOUT
804        )
805        is_accept = pub_accept_event.data.get(constants.SnippetEventParams.IS_ACCEPT, False)
806        if not is_accept:
807            error = pub_accept_event.data[constants.SnippetEventParams.ERROR]
808            asserts.fail(
809                f'{self.publisher} Failed to accept the connection. Error: {error}'
810            )
811
812    def _start_attach(
813        self,
814        ad: android_device.AndroidDevice,
815        is_ranging_enabled: bool,
816    ) -> str:
817        """Starts the attach process on the provided device."""
818        attach_handler = ad.wifi_aware_snippet.wifiAwareAttached(
819            is_ranging_enabled
820        )
821        attach_event = attach_handler.waitAndGet(
822            event_name=AttachCallBackMethodType.ATTACHED,
823            timeout=_DEFAULT_TIMEOUT,
824        )
825        asserts.assert_true(
826            ad.wifi_aware_snippet.wifiAwareIsSessionAttached(attach_event.callback_id),
827            f'{ad} attach succeeded, but Wi-Fi Aware session is still null.'
828        )
829        mac_address = None
830        if is_ranging_enabled:
831            identity_changed_event = attach_handler.waitAndGet(
832                event_name='WifiAwareAttachOnIdentityChanged',
833                timeout=_DEFAULT_TIMEOUT,
834            )
835            mac_address = identity_changed_event.data.get('mac', None)
836            asserts.assert_true(bool(mac_address), 'Mac address should not be empty')
837        ad.log.info('Attach Wi-Fi Aware session succeeded.')
838        return attach_event.callback_id, mac_address
839
840    def _start_publish(
841        self,
842        attach_session_id: str,
843        pub_config: constants.PublishConfig,
844    ) -> callback_event.CallbackEvent:
845        """Starts a publish session on the publisher device."""
846
847        # Start the publishing session and return the handler.
848        publish_handler = self.publisher.wifi_aware_snippet.wifiAwarePublish(
849            attach_session_id,
850            pub_config.to_dict(),
851        )
852
853        # Wait for publish session to start.
854        discovery_event = publish_handler.waitAndGet(
855            event_name=constants.DiscoverySessionCallbackMethodType.DISCOVER_RESULT,
856            timeout=_DEFAULT_TIMEOUT
857        )
858        callback_name = discovery_event.data[_CALLBACK_NAME]
859        asserts.assert_equal(
860            constants.DiscoverySessionCallbackMethodType.PUBLISH_STARTED,
861            callback_name,
862            f'{self.publisher} publish failed, got callback: {callback_name}.',
863        )
864
865        is_session_init = discovery_event.data[_IS_SESSION_INIT]
866        asserts.assert_true(
867            is_session_init,
868            f'{self.publisher} publish succeeded, but null discovery session returned.'
869        )
870        return publish_handler
871
872    def _start_subscribe(
873        self,
874        attach_session_id: str,
875        sub_config: constants.SubscribeConfig,
876    ) -> callback_event.CallbackEvent:
877        """Starts a subscribe session on the subscriber device."""
878
879        # Start the subscription session and return the handler.
880        subscribe_handler = self.subscriber.wifi_aware_snippet.wifiAwareSubscribe(
881            attach_session_id,
882            sub_config.to_dict(),
883        )
884
885        # Wait for subscribe session to start.
886        discovery_event = subscribe_handler.waitAndGet(
887            event_name=constants.DiscoverySessionCallbackMethodType.DISCOVER_RESULT,
888            timeout=_DEFAULT_TIMEOUT
889        )
890        callback_name = discovery_event.data[_CALLBACK_NAME]
891        asserts.assert_equal(
892            constants.DiscoverySessionCallbackMethodType.SUBSCRIBE_STARTED,
893            callback_name,
894            f'{self.subscriber} subscribe failed, got callback: {callback_name}.'
895        )
896        is_session_init = discovery_event.data[_IS_SESSION_INIT]
897        asserts.assert_true(
898            is_session_init,
899            f'{self.subscriber} subscribe succeeded, but null session returned.'
900        )
901        return subscribe_handler
902
903    def _wait_for_discovery(
904        self,
905        sub_aware_session_cb_handler: callback_handler_v2.CallbackHandlerV2,
906        pub_service_specific_info: bytes,
907        is_ranging_enabled: bool,
908    ) -> int:
909        """Waits for discovery of the publisher's service by the subscriber."""
910        event_name = constants.DiscoverySessionCallbackMethodType.SERVICE_DISCOVERED
911        if is_ranging_enabled:
912            event_name = (
913                constants.DiscoverySessionCallbackMethodType.SERVICE_DISCOVERED_WITHIN_RANGE
914            )
915        discover_data = sub_aware_session_cb_handler.waitAndGet(
916            event_name=event_name, timeout=_DEFAULT_TIMEOUT
917        )
918
919        service_info = bytes(
920            discover_data.data[constants.WifiAwareSnippetParams.SERVICE_SPECIFIC_INFO]
921        )
922        asserts.assert_equal(
923            service_info,
924            pub_service_specific_info,
925            f'{self.subscriber} got unexpected service info in discovery'
926            f' callback event "{event_name}".'
927        )
928        match_filters = discover_data.data[
929            constants.WifiAwareSnippetParams.MATCH_FILTER]
930        match_filters = [
931            bytes(filter[constants.WifiAwareSnippetParams.MATCH_FILTER_VALUE])
932            for filter in match_filters
933        ]
934        asserts.assert_equal(
935            match_filters,
936            [constants.WifiAwareTestConstants.MATCH_FILTER_BYTES],
937            f'{self.subscriber} got unexpected match filter data in discovery'
938            f' callback event "{event_name}".'
939        )
940        return discover_data.data[constants.WifiAwareSnippetParams.PEER_ID]
941
942    def _send_msg_through_discovery_session(
943        self,
944        *,
945        sender: android_device.AndroidDevice,
946        sender_aware_session_cb_handler: callback_handler_v2.CallbackHandlerV2,
947        receiver: android_device.AndroidDevice,
948        receiver_aware_session_cb_handler: callback_handler_v2.CallbackHandlerV2,
949        discovery_session: str,
950        peer: int,
951        send_message: str,
952        send_message_id: int,
953    ) -> int:
954        sender.wifi_aware_snippet.wifiAwareSendMessage(
955            discovery_session, peer, send_message_id, send_message
956        )
957        message_send_result = sender_aware_session_cb_handler.waitAndGet(
958            event_name=constants.DiscoverySessionCallbackMethodType.MESSAGE_SEND_RESULT,
959            timeout=_DEFAULT_TIMEOUT,
960        )
961        callback_name = message_send_result.data[
962            constants.DiscoverySessionCallbackParamsType.CALLBACK_NAME
963        ]
964        asserts.assert_equal(
965            callback_name,
966            constants.DiscoverySessionCallbackMethodType.MESSAGE_SEND_SUCCEEDED,
967            f'{sender} failed to send message with an unexpected callback.',
968        )
969        actual_send_message_id = message_send_result.data[
970            constants.DiscoverySessionCallbackParamsType.MESSAGE_ID
971        ]
972        asserts.assert_equal(
973            actual_send_message_id,
974            send_message_id,
975            f'{sender} send message succeeded but message ID mismatched.'
976        )
977        receive_message_event = receiver_aware_session_cb_handler.waitAndGet(
978            event_name=constants.DiscoverySessionCallbackMethodType.MESSAGE_RECEIVED,
979            timeout=_DEFAULT_TIMEOUT,
980        )
981        received_message_raw = receive_message_event.data[
982            constants.WifiAwareSnippetParams.RECEIVED_MESSAGE
983        ]
984        received_message = bytes(received_message_raw).decode('utf-8')
985        asserts.assert_equal(
986            received_message,
987            send_message,
988            f'{receiver} received the message but message content mismatched.'
989        )
990        return receive_message_event.data[constants.WifiAwareSnippetParams.PEER_ID]
991
992    def _request_network(
993        self,
994        ad: android_device.AndroidDevice,
995        discovery_session: str,
996        peer: int,
997        net_work_request_id: str,
998        network_specifier_params: constants.WifiAwareNetworkSpecifier | None = None,
999        is_accept_any_peer: bool = False,
1000    ) -> callback_handler_v2.CallbackHandlerV2:
1001        """Requests and configures a Wi-Fi Aware network connection."""
1002        network_specifier_parcel = (
1003            ad.wifi_aware_snippet.wifiAwareCreateNetworkSpecifier(
1004                discovery_session,
1005                peer,
1006                is_accept_any_peer,
1007                network_specifier_params.to_dict() if network_specifier_params else None,
1008            )
1009        )
1010        network_request_dict = constants.NetworkRequest(
1011            transport_type=_TRANSPORT_TYPE_WIFI_AWARE,
1012            network_specifier_parcel=network_specifier_parcel,
1013        ).to_dict()
1014        ad.log.debug('Requesting Wi-Fi Aware network: %r', network_request_dict)
1015        return ad.wifi_aware_snippet.connectivityRequestNetwork(
1016            net_work_request_id, network_request_dict, _REQUEST_NETWORK_TIMEOUT_MS
1017        )
1018
1019    def _wait_for_network(
1020        self,
1021        ad: android_device.AndroidDevice,
1022        request_network_cb_handler: callback_handler_v2.CallbackHandlerV2,
1023        expected_channel: str | None = None,
1024    ):
1025        """Waits for and verifies the establishment of a Wi-Fi Aware network."""
1026        network_callback_event = request_network_cb_handler.waitAndGet(
1027            event_name=constants.NetworkCbEventName.NETWORK_CALLBACK,
1028            timeout=_DEFAULT_TIMEOUT,
1029        )
1030        callback_name = network_callback_event.data[_CALLBACK_NAME]
1031        if callback_name == constants.NetworkCbName.ON_UNAVAILABLE:
1032            asserts.fail(
1033                f'{ad} failed to request the network, got callback'
1034                f' {callback_name}.'
1035            )
1036        elif callback_name == constants.NetworkCbName.ON_CAPABILITIES_CHANGED:
1037            # `network` is the network whose capabilities have changed.
1038            network = network_callback_event.data[
1039                constants.NetworkCbEventKey.NETWORK]
1040            network_capabilities = network_callback_event.data[
1041                constants.NetworkCbEventKey.NETWORK_CAPABILITIES]
1042            asserts.assert_true(
1043                network and network_capabilities,
1044                f'{ad} received a null Network or NetworkCapabilities!?.'
1045            )
1046            transport_info_class_name = network_callback_event.data[
1047                constants.NetworkCbEventKey.TRANSPORT_INFO_CLASS_NAME]
1048            asserts.assert_equal(
1049                transport_info_class_name,
1050                constants.AWARE_NETWORK_INFO_CLASS_NAME,
1051                f'{ad} network capabilities changes but it is not a WiFi Aware'
1052                ' network.',
1053            )
1054            if expected_channel:
1055                mhz_list = network_callback_event.data[constants.NetworkCbEventKey.CHANNEL_IN_MHZ]
1056                asserts.assert_equal(
1057                    mhz_list,
1058                    [expected_channel],
1059                    f'{ad} Channel freq is not match the request.'
1060                )
1061        else:
1062            asserts.fail(
1063                f'{ad} got unknown request network callback {callback_name}.'
1064            )
1065
1066    def teardown_test(self):
1067        utils.concurrent_exec(
1068            lambda d: d.services.create_output_excerpts_all(self.current_test_info),
1069            param_list=[[ad] for ad in self.ads],
1070            raise_on_exception=True,
1071        )
1072        utils.concurrent_exec(
1073            self._teardown_on_device,
1074            ((self.publisher,), (self.subscriber,)),
1075            max_workers=2,
1076            raise_on_exception=True,
1077        )
1078        self.publisher_mac = None
1079        self.subscriber_mac = None
1080        self.publisher_peer = None
1081        self.subscriber_peer = None
1082        self.publish_session = None
1083        self.subscribe_session = None
1084        self.publisher_attach_session = None
1085        self.subscriber_attach_session = None
1086
1087    def _send_socket_msg(
1088        self,
1089        sender_ad: android_device.AndroidDevice,
1090        receiver_ad: android_device.AndroidDevice,
1091        msg: str,
1092        send_callback_id: str,
1093        receiver_callback_id: str,
1094    ):
1095        """Sends a message from one device to another and verifies receipt."""
1096        is_write_socket = sender_ad.wifi_aware_snippet.connectivityWriteSocket(
1097            send_callback_id, msg
1098        )
1099        asserts.assert_true(
1100            is_write_socket,
1101            f'{sender_ad} Failed to write data to the socket.'
1102        )
1103        sender_ad.log.info('Wrote data to the socket.')
1104        self.publisher.log.info('Server socket accepted the connection.')
1105        # Verify received message
1106        received_message = receiver_ad.wifi_aware_snippet.connectivityReadSocket(
1107            receiver_callback_id, len(msg)
1108        )
1109        asserts.assert_equal(
1110            received_message,
1111            msg,
1112            f'{receiver_ad} received message mismatched.Failure:Expected {msg} but got '
1113            f'{received_message}.'
1114        )
1115        receiver_ad.log.info('Read data from the socket.')
1116
1117    def _skip_if_wifi_rtt_is_not_supported(self):
1118      """Skips this test case if Wi-Fi RTT is not supported on any device."""
1119      asserts.skip_if(
1120          not self.publisher.wifi_aware_snippet.wifiAwareIsRttSupported(),
1121          f'Publisher {self.publisher} does not support Wi-Fi RTT.'
1122      )
1123      asserts.skip_if(
1124          not self.subscriber.wifi_aware_snippet.wifiAwareIsRttSupported(),
1125          f'Subscriber {self.subscriber} does not support Wi-Fi RTT.'
1126      )
1127
1128    def _perform_ranging(
1129        self,
1130        ad: android_device.AndroidDevice,
1131        request: constants.RangingRequest,
1132    ):
1133        """Performs ranging and checks the ranging result."""
1134        ad.log.debug('Starting ranging with request: %s', request)
1135        ranging_cb_handler = ad.wifi_aware_snippet.wifiAwareStartRanging(
1136            request.to_dict()
1137        )
1138        event = ranging_cb_handler.waitAndGet(
1139            event_name=constants.RangingResultCb.EVENT_NAME_ON_RANGING_RESULT,
1140            timeout=_DEFAULT_TIMEOUT,
1141        )
1142
1143        callback_name = event.data.get(
1144            constants.RangingResultCb.DATA_KEY_CALLBACK_NAME, None
1145        )
1146        asserts.assert_equal(
1147            callback_name,
1148            constants.RangingResultCb.CB_METHOD_ON_RANGING_RESULT,
1149            'Ranging failed: got unexpected callback.',
1150        )
1151
1152        results = event.data.get(
1153            constants.RangingResultCb.DATA_KEY_RESULTS, None
1154        )
1155        asserts.assert_true(
1156            results is not None and len(results) == 1,
1157            'Ranging got invalid results: null, empty, or wrong length.'
1158        )
1159
1160        status_code = results[0].get(
1161            constants.RangingResultCb.DATA_KEY_RESULT_STATUS, None
1162        )
1163        asserts.assert_equal(
1164            status_code,
1165            constants.RangingResultStatusCode.SUCCESS,
1166            'Ranging peer failed: invalid result status code.',
1167        )
1168
1169        distance_mm = results[0].get(
1170            constants.RangingResultCb.DATA_KEY_RESULT_DISTANCE_MM, None
1171        )
1172        asserts.assert_true(
1173            (
1174                distance_mm is not None
1175                and distance_mm <= _LARGE_ENOUGH_DISTANCE_MM
1176            ),
1177            'Ranging peer failed: invalid distance in ranging result.',
1178        )
1179        rssi = results[0].get(
1180            constants.RangingResultCb.DATA_KEY_RESULT_RSSI, None
1181        )
1182        asserts.assert_true(
1183            rssi is not None and rssi >= _MIN_RSSI,
1184            'Ranging peer failed: invalid rssi in ranging result.',
1185        )
1186
1187        peer_id = results[0].get(
1188            constants.RangingResultCb.DATA_KEY_PEER_ID, None
1189        )
1190        if peer_id is not None:
1191            msg = 'Ranging peer failed: invalid peer ID in ranging result.'
1192            asserts.assert_in(peer_id, request.peer_ids, msg)
1193
1194        peer_mac = results[0].get(constants.RangingResultCb.DATA_KEY_MAC, None)
1195        if peer_mac is not None:
1196            msg = (
1197                'Ranging peer failed: invalid peer MAC address in ranging '
1198                'result.'
1199            )
1200            asserts.assert_in(peer_mac, request.peer_mac_addresses, msg)
1201
1202    def _teardown_on_device(self, ad: android_device.AndroidDevice) -> None:
1203        """Releases resources and sessions after each test."""
1204        ad.wifi_aware_snippet.connectivityReleaseAllSockets()
1205        ad.wifi_aware_snippet.wifiAwareCloseAllWifiAwareSession()
1206
1207    def on_fail(self, record: records.TestResult) -> None:
1208        logging.info('Collecting bugreports...')
1209        android_device.take_bug_reports(self.ads, destination=self.current_test_info.output_path)
1210
1211
1212if __name__ == '__main__':
1213    # Take test args
1214    if '--' in sys.argv:
1215        index = sys.argv.index('--')
1216        sys.argv = sys.argv[:1] + sys.argv[index + 1:]
1217
1218    test_runner.main()
1219