cases: move common test routines to avatar

4 files changed, 76 insertions, 52 deletions

diff --git a/avatar/pandora.py b/avatar/pandora.py
new file mode 100644
index 0000000..31695ee
--- /dev/null
+++ b/avatar/pandora.py
@@ -0,0 +1,67 @@
+# Copyright 2023 Google LLC
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     https://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import asyncio
+
+from avatar import BumblePandoraDevice
+from avatar import PandoraDevice
+from bumble.device import Connection as BumbleConnection
+from mobly.asserts import assert_equal  # type: ignore
+from mobly.asserts import assert_is_not_none  # type: ignore
+from pandora._utils import AioStream
+from pandora.host_pb2 import AdvertiseResponse
+from pandora.host_pb2 import Connection
+from pandora.host_pb2 import OwnAddressType
+from pandora.host_pb2 import ScanningResponse
+from typing import Optional, Tuple
+
+
+def get_raw_connection_handle(device: PandoraDevice, connection: Connection) -> int:
+    assert isinstance(device, BumblePandoraDevice)
+    return int.from_bytes(connection.cookie.value, 'big')
+
+
+def get_raw_connection(device: PandoraDevice, connection: Connection) -> Optional[BumbleConnection]:
+    assert isinstance(device, BumblePandoraDevice)
+    return device.device.lookup_connection(get_raw_connection_handle(device, connection))
+
+
+async def connect(initiator: PandoraDevice, acceptor: PandoraDevice) -> Tuple[Connection, Connection]:
+    init_res, wait_res = await asyncio.gather(
+        initiator.aio.host.Connect(address=acceptor.address),
+        acceptor.aio.host.WaitConnection(address=initiator.address),
+    )
+    assert_equal(init_res.result_variant(), 'connection')
+    assert_equal(wait_res.result_variant(), 'connection')
+    assert init_res.connection is not None and wait_res.connection is not None
+    return init_res.connection, wait_res.connection
+
+
+async def connect_le(
+    initiator: PandoraDevice,
+    acceptor: AioStream[AdvertiseResponse],
+    scan: ScanningResponse,
+    own_address_type: OwnAddressType,
+    cancel_advertisement: bool = True,
+) -> Tuple[Connection, Connection]:
+    (init_res, wait_res) = await asyncio.gather(
+        initiator.aio.host.ConnectLE(own_address_type=own_address_type, **scan.address_asdict()),
+        anext(aiter(acceptor)),  # pytype: disable=name-error
+    )
+    if cancel_advertisement:
+        acceptor.cancel()
+    assert_equal(init_res.result_variant(), 'connection')
+    assert_is_not_none(init_res.connection)
+    assert init_res.connection
+    return init_res.connection, wait_res.connection
diff --git a/cases/le_security_test.py b/cases/le_security_test.py
index 3cdc224..b44465f 100644
--- a/cases/le_security_test.py
+++ b/cases/le_security_test.py
@@ -20,6 +20,7 @@ import logging
 from avatar import BumblePandoraDevice
 from avatar import PandoraDevice
 from avatar import PandoraDevices
+from avatar import pandora
 from bumble.pairing import PairingConfig
 from bumble.pairing import PairingDelegate
 from mobly import base_test
@@ -40,7 +41,6 @@ from pandora.security_pb2 import PairingEventAnswer
 from pandora.security_pb2 import SecureResponse
 from pandora.security_pb2 import WaitSecurityResponse
 from typing import Any, Literal, Optional, Tuple, Union
-from utils import make_bredr_connection
 
 
 class LeSecurityTest(base_test.BaseTestClass):  # type: ignore[misc]
@@ -220,24 +220,13 @@ class LeSecurityTest(base_test.BaseTestClass):  # type: ignore[misc]
 
                 # Initiator - Scan and fetch the address
                 scan = initiator.aio.host.Scan(own_address_type=initiator_addr_type)
-                acceptor_addr = await anext(
+                acceptor_scan = await anext(
                     (x async for x in scan if b'pause cafe' in x.data.manufacturer_specific_data)
                 )  # pytype: disable=name-error
                 scan.cancel()
 
                 # Initiator - LE connect
-                init_res, wait_res = await asyncio.gather(
-                    initiator.aio.host.ConnectLE(
-                        own_address_type=initiator_addr_type, **acceptor_addr.address_asdict()
-                    ),
-                    anext(aiter(advertisement)),  # pytype: disable=name-error
-                )
-
-                advertisement.cancel()
-                assert_equal(init_res.result_variant(), 'connection')
-
-                assert init_res.connection is not None and wait_res.connection is not None
-                return init_res.connection, wait_res.connection
+                return await pandora.connect_le(initiator, advertisement, acceptor_scan, initiator_addr_type)
 
             # Make LE connection.
             if connect == 'incoming_connection':
@@ -379,7 +368,7 @@ class LeSecurityTest(base_test.BaseTestClass):  # type: ignore[misc]
                     assert_equal(wait_security.result_variant(), 'success')
                     if 'lk' in key_distribution:
                         # Make a Classic connection
-                        ref_dut_classic, _dut_ref_clsssic = await make_bredr_connection(self.ref, self.dut)
+                        ref_dut_classic, _dut_ref_clsssic = await pandora.connect(self.ref, self.dut)
                         # Try to encrypt Classic connection
                         ref_dut_secure = await self.ref.aio.security.Secure(ref_dut_classic, classic=LEVEL2)
                         assert_equal(ref_dut_secure.result_variant(), 'success')
diff --git a/cases/security_test.py b/cases/security_test.py
index a5e3c4d..85356fb 100644
--- a/cases/security_test.py
+++ b/cases/security_test.py
@@ -20,6 +20,7 @@ import logging
 from avatar import BumblePandoraDevice
 from avatar import PandoraDevice
 from avatar import PandoraDevices
+from avatar import pandora
 from bumble.hci import HCI_CENTRAL_ROLE
 from bumble.hci import HCI_PERIPHERAL_ROLE
 from bumble.hci import HCI_Write_Default_Link_Policy_Settings_Command
@@ -43,7 +44,6 @@ from pandora.security_pb2 import PairingEventAnswer
 from pandora.security_pb2 import SecureResponse
 from pandora.security_pb2 import WaitSecurityResponse
 from typing import Any, List, Literal, Optional, Tuple, Union
-from utils import make_bredr_connection
 
 DEFAULT_SMP_KEY_DISTRIBUTION = (
     PairingDelegate.KeyDistribution.DISTRIBUTE_ENCRYPTION_KEY
@@ -254,16 +254,16 @@ class SecurityTest(base_test.BaseTestClass):  # type: ignore[misc]
 
             # Make classic connection.
             if connect == 'incoming_connection':
-                ref_dut, dut_ref = await make_bredr_connection(initiator=self.ref, acceptor=self.dut)
+                ref_dut, dut_ref = await pandora.connect(initiator=self.ref, acceptor=self.dut)
             else:
-                dut_ref, ref_dut = await make_bredr_connection(initiator=self.dut, acceptor=self.ref)
+                dut_ref, ref_dut = await pandora.connect(initiator=self.dut, acceptor=self.ref)
 
             # Retrieve Bumble connection
             if isinstance(self.dut, BumblePandoraDevice):
-                dut_ref_bumble = self.dut.device.lookup_connection(int.from_bytes(dut_ref.cookie.value, 'big'))  # type: ignore
+                dut_ref_bumble = pandora.get_raw_connection(self.dut, dut_ref)
             # Role switch.
             if isinstance(self.ref, BumblePandoraDevice):
-                ref_dut_bumble = self.ref.device.lookup_connection(int.from_bytes(ref_dut.cookie.value, 'big'))  # type: ignore
+                ref_dut_bumble = pandora.get_raw_connection(self.ref, ref_dut)
                 if ref_dut_bumble is not None:
                     role = {
                         'against_central': HCI_CENTRAL_ROLE,
diff --git a/cases/utils.py b/cases/utils.py
deleted file mode 100644
index 57e2ade..0000000
--- a/cases/utils.py
+++ /dev/null
@@ -1,32 +0,0 @@
-# Copyright 2023 Google LLC
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-#     https://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-
-import asyncio
-
-from avatar import PandoraDevice
-from mobly.asserts import assert_equal  # type: ignore
-from pandora.host_pb2 import Connection
-from typing import Tuple
-
-
-# Make classic connection task.
-async def make_bredr_connection(initiator: PandoraDevice, acceptor: PandoraDevice) -> Tuple[Connection, Connection]:
-    init_res, wait_res = await asyncio.gather(
-        initiator.aio.host.Connect(address=acceptor.address),
-        acceptor.aio.host.WaitConnection(address=initiator.address),
-    )
-    assert_equal(init_res.result_variant(), 'connection')
-    assert_equal(wait_res.result_variant(), 'connection')
-    assert init_res.connection is not None and wait_res.connection is not None
-    return init_res.connection, wait_res.connection