async def test_remove_companion_by_udid(self) -> None: companion_manager = CompanionManager(logger=mock.MagicMock()) add_companion(companion_manager, TEST_COMPANION) self.assertEqual(len(companion_manager._udid_companion_map), 1) self.assertEqual(len(companion_manager._udid_target_map), 1) result = companion_manager.remove_companion(TEST_COMPANION.udid) self.assertEqual(result, TEST_COMPANION) self.assertEqual(len(companion_manager._udid_companion_map), 0) self.assertEqual(len(companion_manager._udid_target_map), 0)
async def test_remove_companion_by_address(self) -> None: companion_manager = CompanionManager(companion_path=None, logger=mock.MagicMock()) add_companion(companion_manager, TEST_COMPANION) self.assertEqual(len(companion_manager._udid_companion_map), 1) self.assertEqual(len(companion_manager._udid_target_map), 1) result = companion_manager.remove_companion( Address(host=TEST_COMPANION.host, port=TEST_COMPANION.port)) self.assertEqual(result, TEST_COMPANION) self.assertEqual(len(companion_manager._udid_companion_map), 0) self.assertEqual(len(companion_manager._udid_target_map), 0)