async def test_add_multiple(self) -> None: async for manager in self._managers(): companion_a = CompanionInfo(udid="a", host="ahost", port=123, is_local=False) replaced = await manager.add_companion(companion_a) self.assertIsNone(replaced) companions = await manager.get_companions() self.assertEqual(companions, [companion_a]) companion_b = CompanionInfo(udid="b", host="bhost", port=123, is_local=False) replaced = await manager.add_companion(companion_b) self.assertIsNone(replaced) companions = await manager.get_companions() self.assertEqual(companions, [companion_a, companion_b]) companion_c = CompanionInfo(udid="c", host="chost", port=123, is_local=False) replaced = await manager.add_companion(companion_c) self.assertIsNone(replaced) companions = await manager.get_companions() self.assertEqual(companions, [companion_a, companion_b, companion_c]) removed = await manager.remove_companion( Address(host=companion_b.host, port=companion_b.port)) companions = await manager.get_companions() self.assertEqual(companions, [companion_a, companion_c]) self.assertEqual(removed, [companion_b]) removed = await manager.remove_companion("a") companions = await manager.get_companions() self.assertEqual(companions, [companion_c])
async def _spawn_companion(self, target_udid: str) -> Optional[CompanionInfo]: companion_spawner = self._companion_spawner if companion_spawner is None: return None local_target_available = ( await self._local_targets_manager.is_local_target_available( target_udid=target_udid)) if local_target_available or target_udid == "mac": self._logger.info( f"will attempt to spawn a companion for {target_udid}") port = await companion_spawner.spawn_companion( target_udid=target_udid) if port: self._logger.info(f"spawned a companion for {target_udid}") host = "localhost" companion_info = CompanionInfo( address=TCPAddress(host=host, port=port), udid=target_udid, is_local=True, ) await self._direct_companion_manager.add_companion( companion_info) return companion_info return None
async def connect( self, destination: ConnectionDestination, metadata: Optional[Dict[str, str]] = None, ) -> CompanionInfo: self.logger.debug(f"Connecting directly to {destination} with meta {metadata}") if isinstance(destination, Address): channel = Channel( destination.host, destination.port, loop=asyncio.get_event_loop() ) stub = CompanionServiceStub(channel=channel) with tempfile.NamedTemporaryFile(mode="w+b") as f: response = await stub.connect( ConnectRequest( destination=destination_to_grpc(destination), metadata=metadata, local_file_path=f.name, ) ) companion = CompanionInfo( udid=response.companion.udid, host=destination.host, port=destination.port, is_local=response.companion.is_local, ) self.logger.debug(f"Connected directly to {companion}") self.direct_companion_manager.add_companion(companion) channel.close() return companion else: companion = await self.spawn_companion(target_udid=destination) if companion: return companion else: raise IdbException(f"can't find target for udid {destination}")
async def get_stub(self) -> CompanionServiceStub: channel: Optional[Channel] = None try: try: self.companion_info = self.direct_companion_manager.get_companion_info( target_udid=self.target_udid ) except IdbException as e: # will try to spawn a companion if on mac. if platform.system() == "Darwin" and self.target_udid: self.logger.info( f"will attempt to spawn a companion for {self.target_udid}" ) port = await self.companion_spawner.spawn_companion( target_udid=self.target_udid ) host = "localhost" self.companion_info = CompanionInfo( host=host, port=port, udid=self.target_udid, is_local=True ) self.direct_companion_manager.add_companion(self.companion_info) else: raise e self.logger.info(f"using companion {self.companion_info}") channel = Channel( self.companion_info.host, self.companion_info.port, loop=asyncio.get_event_loop(), ) yield CompanionServiceStub(channel=channel) finally: if channel: channel.close()
def companion_to_py(companion: GrpcCompanionInfo) -> CompanionInfo: return CompanionInfo( udid=companion.udid, host=companion.host, port=companion.port, is_local=companion.is_local, )
async def connect( self, destination: ConnectionDestination, metadata: Optional[Dict[str, str]] = None, ) -> CompanionInfo: self._logger.debug( f"Connecting directly to {destination} with meta {metadata}") if isinstance(destination, TCPAddress) or isinstance( destination, DomainSocketAddress): async with IdbClient.build(address=destination, is_local=False, logger=self._logger) as client: with tempfile.NamedTemporaryFile(mode="w+b") as f: response = await client.stub.connect( ConnectRequest(metadata=metadata, local_file_path=f.name)) companion = CompanionInfo( address=destination, udid=response.companion.udid, is_local=response.companion.is_local, ) self._logger.debug(f"Connected directly to {companion}") await self._direct_companion_manager.add_companion(companion) return companion else: companion = await self._spawn_companion(target_udid=destination) if companion: return companion else: raise IdbException(f"can't find target for udid {destination}")
def companion_to_py(companion: GrpcCompanionInfo) -> CompanionInfo: return CompanionInfo( udid=companion.udid, host=companion.host if ".local" not in companion.host else "localhost", port=companion.grpc_port, is_local=companion.is_local, )
async def test_selects_by_udid(self) -> None: async for manager in self._managers(): # Add two companions companion_a = CompanionInfo(udid="a", host="ahost", port=123, is_local=False) await manager.add_companion(companion_a) companion_b = CompanionInfo(udid="b", host="bhost", port=123, is_local=False) await manager.add_companion(companion_b) self.assertEqual(companion_a, await manager.get_companion_info(target_udid="a")) self.assertEqual(companion_b, await manager.get_companion_info(target_udid="b"))
async def test_add_companion_with_host_and_port_adds_companion( self) -> None: companion_manager = CompanionManager(logger=mock.MagicMock()) companion_manager.add_companion( CompanionInfo(udid="asdasda", host="foohost", port=123, is_local=False)) assert companion_manager._udid_companion_map["asdasda"]
async def describe(self) -> TargetDescription: response = await self.stub.describe(TargetDescriptionRequest()) target = response.target_description return target_to_py( target=target, companion=CompanionInfo(address=self.address, udid=target.udid, is_local=self.is_local), )
async def test_selects_when_no_udid_single_companion(self) -> None: async for manager in self._managers(): companion = CompanionInfo(udid="a", host="ahost", port=123, is_local=False) await manager.add_companion(companion) self.assertEqual( companion, await manager.get_companion_info(target_udid=None))
async def test_ambiguous_companions(self) -> None: async for manager in self._managers(): companion_a = CompanionInfo(udid="a", host="ahost", port=123, is_local=False) await manager.add_companion(companion_a) companions = await manager.get_companions() self.assertEqual(companions, [companion_a]) companion_b = CompanionInfo(udid="b", host="ahost", port=123, is_local=False) await manager.add_companion(companion_b) companions = await manager.get_companions() self.assertEqual(companions, [companion_a, companion_b]) with self.assertRaises(IdbException): await manager.get_companion_info(target_udid=None)
def companion_to_py(companion: GrpcCompanionInfo, address: Address) -> CompanionInfo: metadata = companion.metadata return CompanionInfo( address=address, udid=companion.udid, is_local=companion.is_local, metadata=(json.loads(metadata.decode()) if len(metadata) else {}), )
async def test_replace_companion(self) -> None: async for manager in self._managers(): companion_first = CompanionInfo(udid="a", host="ahost", port=123, is_local=False) replaced = await manager.add_companion(companion_first) self.assertIsNone(replaced) companions = await manager.get_companions() self.assertEqual(companions, [companion_first]) companion_second = CompanionInfo(udid="a", host="anotherhost", port=321, is_local=False) replaced = await manager.add_companion(companion_second) self.assertEqual(replaced, companion_first) companions = await manager.get_companions() self.assertEqual(companions, [companion_second])
def companion_to_py(companion: GrpcCompanionInfo, address: Address, is_local: Optional[bool]) -> CompanionInfo: metadata = companion.metadata return CompanionInfo( address=address, udid=companion.udid, is_local=(is_local if is_local is not None else companion.is_local), metadata=(json.loads(metadata.decode()) if len(metadata) else {}), )
def json_to_companion_info(data: List[Dict[str, Any]]) -> List[CompanionInfo]: return [ CompanionInfo( udid=item["udid"], address=(TCPAddress(host=item["host"], port=item["port"]) if "host" in item else DomainSocketAddress(path=item["path"])), is_local=item["is_local"], pid=item.get("pid"), ) for item in data ]
async def describe(self, fetch_diagnostics: bool = False) -> TargetDescription: response = await self.stub.describe( TargetDescriptionRequest(fetch_diagnostics=fetch_diagnostics) ) target = response.target_description return target_to_py( target=target, companion=CompanionInfo( address=self.address, udid=target.udid, is_local=self.is_local ), )
def json_to_companion_info(data: List[Dict[str, Any]]) -> List[CompanionInfo]: companion_list = [] for item in data: companion_list.append( CompanionInfo( udid=item["udid"], host=item["host"], port=item["port"], is_local=item["is_local"], )) return companion_list
async def test_add_then_clear(self) -> None: async for manager in self._managers(): companion = CompanionInfo(udid="asdasda", host="foohost", port=123, is_local=False) await manager.add_companion(companion) companions = await manager.get_companions() self.assertEqual(companions, [companion]) await manager.clear() companions = await manager.get_companions() self.assertEqual(companions, [])
async def test_clear(self) -> None: with tempfile.NamedTemporaryFile() as f: companion_manager = DirectCompanionManager( logger=mock.MagicMock(), state_file_path=f.name ) companion = CompanionInfo( udid="asdasda", host="foohost", port=123, is_local=False ) companion_manager.add_companion(companion) companion_manager.clear() companions = companion_manager.load_companions() self.assertEqual(companions, [])
async def test_ambiguity_when_no_udid_multiple_companions(self) -> None: async for manager in self._managers(): companion_a = CompanionInfo( udid="a", address=TCPAddress(host="ahost", port=123), is_local=False ) replaced = await manager.add_companion(companion_a) self.assertIsNone(replaced) companions = await manager.get_companions() self.assertEqual(companions, [companion_a]) companion_b = CompanionInfo( udid="b", address=TCPAddress(host="ahost", port=123), is_local=False ) replaced = await manager.add_companion(companion_b) self.assertIsNone(replaced) companions = await manager.get_companions() self.assertEqual(companions, [companion_a, companion_b]) with self.assertRaises(IdbException) as cm: await manager.get_companion_info(target_udid=None) self.assertIn( "No UDID provided and there's multiple companions", str(cm.exception) )
async def test_replace_companion(self) -> None: async for manager in self._managers(): companion_first = CompanionInfo( udid="a", address=TCPAddress(host="ahost", port=123), is_local=False, pid=None, ) replaced = await manager.add_companion(companion_first) self.assertIsNone(replaced) companions = await manager.get_companions() self.assertEqual(companions, [companion_first]) companion_second = CompanionInfo( udid="a", address=DomainSocketAddress(path="/some/path"), is_local=False, pid=123, ) replaced = await manager.add_companion(companion_second) self.assertEqual(replaced, companion_first) companions = await manager.get_companions() self.assertEqual(companions, [companion_second])
async def test_get_companions(self) -> None: with tempfile.NamedTemporaryFile() as f: companion_manager = DirectCompanionManager( logger=mock.MagicMock(), state_file_path=f.name ) companion = CompanionInfo( udid="asdasda", host="foohost", port=123, is_local=False ) with open(f.name, "w") as f: json.dump(json_data_companions([companion]), f) companions = companion_manager._load() read_companion: CompanionInfo = companions[0] self.assertEqual(companion, read_companion)
async def test_add_then_remove_companion_by_udid(self) -> None: async for manager in self._managers(): companion = CompanionInfo(udid="asdasda", host="foohost", port=123, is_local=False) await manager.add_companion(companion) companions = await manager.get_companions() self.assertEqual(companions, [companion]) removed = await manager.remove_companion("asdasda") companions = await manager.get_companions() self.assertEqual(companions, []) self.assertEqual(removed, [companion])
async def test_add_then_clear(self) -> None: async for manager in self._managers(): first = CompanionInfo( udid="asdasda", address=TCPAddress(host="foohost", port=123), is_local=False, pid=None, ) second = CompanionInfo( udid="fooo", address=DomainSocketAddress(path="/bar/bar"), is_local=False, pid=324, ) await manager.add_companion(first) await manager.add_companion(second) companions = await manager.get_companions() self.assertEqual(companions, [first, second]) cleared = await manager.clear() companions = await manager.get_companions() self.assertEqual(companions, []) self.assertEqual(cleared, [first, second])
async def test_add_companion(self) -> None: with tempfile.NamedTemporaryFile() as f: companion_manager = DirectCompanionManager( logger=mock.MagicMock(), state_file_path=f.name ) companion = CompanionInfo( udid="asdasda", host="foohost", port=123, is_local=False ) companion_manager.add_companion(companion) data = json.load(f) companions = json_to_companion_info(data) read_companion: CompanionInfo = companions[0] self.assertEqual(companion, read_companion)
async def test_remove_companion_with_udid(self) -> None: with tempfile.NamedTemporaryFile() as f: companion_manager = DirectCompanionManager(logger=mock.MagicMock(), state_file_path=f.name) companion = CompanionInfo(udid="asdasda", host="foohost", port=123, is_local=False) with open(f.name, "w") as f: json.dump(json_data_companions([companion]), f) companion_manager.remove_companion( Address(host=companion.host, port=companion.port)) companions = companion_manager._load() self.assertEqual(companions, [])
async def describe(self, fetch_diagnostics: bool = False) -> TargetDescription: response = await self.stub.describe( TargetDescriptionRequest(fetch_diagnostics=fetch_diagnostics)) target = response.target_description return target_to_py( target=target, # Use the local understanding of the companion instead of the remote's. companion=CompanionInfo(address=self.address, udid=target.udid, is_local=self.is_local), # Extract the companion metadata from the response. metadata=response.companion.metadata, )
async def test_get_default_companion(self) -> None: companion_manager = CompanionManager(logger=mock.MagicMock()) self.assertFalse(companion_manager.has_default_companion()) add_companion(companion_manager, TEST_COMPANION) self.assertTrue(companion_manager.has_default_companion()) self.assertEqual(companion_manager.get_default_companion(), TEST_COMPANION) companion_manager.add_companion( CompanionInfo( udid="someOtherUdid", host=TEST_COMPANION.host, port=TEST_COMPANION.port, is_local=TEST_COMPANION.is_local, )) self.assertFalse(companion_manager.has_default_companion())
async def test_add_then_remove_companion_by_address(self) -> None: async for manager in self._managers(): companion = CompanionInfo(udid="asdasda", host="foohost", port=123, is_local=False) replaced = await manager.add_companion(companion) self.assertIsNone(replaced) companions = await manager.get_companions() self.assertEqual(companions, [companion]) removed = await manager.remove_companion( Address(host=companion.host, port=companion.port)) companions = await manager.get_companions() self.assertEqual(companions, []) self.assertEqual(removed, [companion])