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])
