1# Copyright 2022 Google LLC
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#     https://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
16"""Interface for controller-specific Pandora server management."""
17
18import asyncio
19import avatar.aio
20import grpc
21import grpc.aio
22import portpicker
23import threading
24import types
25
26from avatar.controllers import bumble_device
27from avatar.controllers import pandora_device
28from avatar.pandora_client import BumblePandoraClient
29from avatar.pandora_client import PandoraClient
30from bumble import pandora as bumble_server
31from bumble.pandora.device import PandoraDevice as BumblePandoraDevice
32from contextlib import suppress
33from mobly.controllers import android_device
34from mobly.controllers.android_device import AndroidDevice
35from typing import Generic, Optional, TypeVar
36
37ANDROID_SERVER_PACKAGE = 'com.android.pandora'
38ANDROID_SERVER_GRPC_PORT = 8999
39
40
41# Generic type for `PandoraServer`.
42TDevice = TypeVar('TDevice')
43
44
45class PandoraServer(Generic[TDevice]):
46    """Abstract interface to manage the Pandora gRPC server on the device."""
47
48    MOBLY_CONTROLLER_MODULE: types.ModuleType = pandora_device
49
50    device: TDevice
51
52    def __init__(self, device: TDevice) -> None:
53        """Creates a PandoraServer.
54
55        Args:
56            device: A Mobly controller instance.
57        """
58        self.device = device
59
60    def start(self) -> PandoraClient:
61        """Sets up and starts the Pandora server on the device."""
62        assert isinstance(self.device, PandoraClient)
63        return self.device
64
65    def stop(self) -> None:
66        """Stops and cleans up the Pandora server on the device."""
67
68
69class BumblePandoraServer(PandoraServer[BumblePandoraDevice]):
70    """Manages the Pandora gRPC server on a BumbleDevice."""
71
72    MOBLY_CONTROLLER_MODULE = bumble_device
73
74    _task: Optional[asyncio.Task[None]] = None
75
76    def start(self) -> BumblePandoraClient:
77        """Sets up and starts the Pandora server on the Bumble device."""
78        assert self._task is None
79
80        # set the event loop to make sure the gRPC server use the avatar one.
81        asyncio.set_event_loop(avatar.aio.loop)
82
83        # create gRPC server & port.
84        server = grpc.aio.server()
85        port = server.add_insecure_port(f'localhost:{0}')
86
87        config = bumble_server.Config()
88        self._task = avatar.aio.loop.create_task(
89            bumble_server.serve(self.device, config=config, grpc_server=server, port=port)
90        )
91
92        return BumblePandoraClient(f'localhost:{port}', self.device, config)
93
94    def stop(self) -> None:
95        """Stops and cleans up the Pandora server on the Bumble device."""
96
97        async def server_stop() -> None:
98            assert self._task is not None
99            if not self._task.done():
100                self._task.cancel()
101                with suppress(asyncio.CancelledError):
102                    await self._task
103            self._task = None
104
105        avatar.aio.run_until_complete(server_stop())
106
107
108class AndroidPandoraServer(PandoraServer[AndroidDevice]):
109    """Manages the Pandora gRPC server on an AndroidDevice."""
110
111    MOBLY_CONTROLLER_MODULE = android_device
112
113    _instrumentation: Optional[threading.Thread] = None
114    _port: int
115
116    def start(self) -> PandoraClient:
117        """Sets up and starts the Pandora server on the Android device."""
118        assert self._instrumentation is None
119
120        # start Pandora Android gRPC server.
121        self._port = portpicker.pick_unused_port()  # type: ignore
122        self._instrumentation = threading.Thread(
123            target=lambda: self.device.adb._exec_adb_cmd(  # type: ignore
124                'shell',
125                f'am instrument --no-hidden-api-checks -w {ANDROID_SERVER_PACKAGE}/.Main',
126                shell=False,
127                timeout=None,
128                stderr=None,
129            )
130        )
131
132        self._instrumentation.start()
133        self.device.adb.forward([f'tcp:{self._port}', f'tcp:{ANDROID_SERVER_GRPC_PORT}'])  # type: ignore
134
135        return PandoraClient(f'localhost:{self._port}', 'android')
136
137    def stop(self) -> None:
138        """Stops and cleans up the Pandora server on the Android device."""
139        assert self._instrumentation is not None
140
141        # Stop Pandora Android gRPC server.
142        self.device.adb._exec_adb_cmd(  # type: ignore
143            'shell', f'am force-stop {ANDROID_SERVER_PACKAGE}', shell=False, timeout=None, stderr=None
144        )
145
146        self.device.adb.forward(['--remove', f'tcp:{self._port}'])  # type: ignore
147        self._instrumentation.join()
148        self._instrumentation = None
149