← back to Exo
fix test runtime warning
8d24df2b4b52346b3b87cba247a5b8fefac39b18 · 2024-10-24 17:03:50 +0700 · Ian Paul
Files touched
M exo/networking/manual/test_manual_discovery.py
Diff
commit 8d24df2b4b52346b3b87cba247a5b8fefac39b18
Author: Ian Paul <iandouglaspaul@gmail.com>
Date: Thu Oct 24 17:03:50 2024 +0700
fix test runtime warning
---
exo/networking/manual/test_manual_discovery.py | 301 ++++++++++++++-----------
1 file changed, 169 insertions(+), 132 deletions(-)
diff --git a/exo/networking/manual/test_manual_discovery.py b/exo/networking/manual/test_manual_discovery.py
index 9efe7d1b..8af24ce9 100644
--- a/exo/networking/manual/test_manual_discovery.py
+++ b/exo/networking/manual/test_manual_discovery.py
@@ -12,146 +12,183 @@ root_path = "./exo/networking/manual/test_data/test_config.json"
class TestSingleNodeManualDiscovery(unittest.IsolatedAsyncioTestCase):
- async def asyncSetUp(self):
- self.peer1 = mock.AsyncMock()
- self.peer1.connect = mock.AsyncMock()
- self.discovery1 = ManualDiscovery(root_path, "node1", create_peer_handle=lambda peer_id, address, description, device_capabilities: self.peer1)
- _ = self.discovery1.start()
-
- async def asyncTearDown(self):
- await self.discovery1.stop()
-
- async def test_discovery(self):
- peers1 = await self.discovery1.discover_peers(wait_for_peers=0)
- assert len(peers1) == 0
-
- self.peer1.connect.assert_not_called()
+ async def asyncSetUp(self):
+ self.peer1 = mock.AsyncMock(spec=Node)
+ self.peer1.connect = mock.AsyncMock()
+ self.server1 = GRPCServer(self.peer1, "localhost", 8000)
+ await self.server1.start()
+ self.discovery1 = ManualDiscovery(
+ root_path,
+ "node1",
+ create_peer_handle=lambda peer_id, address, device_capabilities: GRPCPeerHandle(
+ peer_id, address, device_capabilities
+ ),
+ )
+ await self.discovery1.start()
+
+ async def asyncTearDown(self):
+ await self.discovery1.stop()
+ await self.server1.stop()
+
+ async def test_discovery(self):
+ peers1 = await self.discovery1.discover_peers(wait_for_peers=0)
+ self.assertEqual(len(peers1), 0)
+
+ self.peer1.connect.assert_not_called()
class TestManualDiscovery(unittest.IsolatedAsyncioTestCase):
- async def asyncSetUp(self):
- self.peer1 = mock.AsyncMock()
- self.peer2 = mock.AsyncMock()
- self.peer1.connect = mock.AsyncMock()
- self.peer2.connect = mock.AsyncMock()
- self.discovery1 = ManualDiscovery(root_path, "node1", create_peer_handle=lambda peer_id, address, description, device_capabilities: self.peer1)
- self.discovery2 = ManualDiscovery(root_path, "node2", create_peer_handle=lambda peer_id, address, description, device_capabilities: self.peer2)
- await self.discovery1.start()
- await self.discovery2.start()
-
- async def asyncTearDown(self):
- await self.discovery1.stop()
- await self.discovery2.stop()
-
- async def test_discovery(self):
- peers1 = await self.discovery1.discover_peers(wait_for_peers=1)
- self.assertEqual(len(peers1), 1)
- peers2 = await self.discovery2.discover_peers(wait_for_peers=1)
- self.assertEqual(len(peers2), 1)
-
- # connect has to be explicitly called after discovery
- self.peer1.connect.assert_not_called()
- self.peer2.connect.assert_not_called()
+ async def asyncSetUp(self):
+ self.peer1 = mock.AsyncMock()
+ self.peer2 = mock.AsyncMock()
+ self.peer1.connect = mock.AsyncMock()
+ self.peer2.connect = mock.AsyncMock()
+ self.discovery1 = ManualDiscovery(
+ root_path,
+ "node1",
+ create_peer_handle=lambda peer_id, address, description, device_capabilities: self.peer1,
+ )
+ self.discovery2 = ManualDiscovery(
+ root_path,
+ "node2",
+ create_peer_handle=lambda peer_id, address, description, device_capabilities: self.peer2,
+ )
+ await self.discovery1.start()
+ await self.discovery2.start()
+
+ async def asyncTearDown(self):
+ await self.discovery1.stop()
+ await self.discovery2.stop()
+
+ async def test_discovery(self):
+ peers1 = await self.discovery1.discover_peers(wait_for_peers=1)
+ self.assertEqual(len(peers1), 1)
+ peers2 = await self.discovery2.discover_peers(wait_for_peers=1)
+ self.assertEqual(len(peers2), 1)
+
+ # connect has to be explicitly called after discovery
+ self.peer1.connect.assert_not_called()
+ self.peer2.connect.assert_not_called()
class TestManualDiscoveryWithGRPCPeerHandle(unittest.IsolatedAsyncioTestCase):
- async def asyncSetUp(self):
- config = NetworkTopology.from_path(root_path)
-
- self.node1 = mock.AsyncMock(spec=Node)
- self.node2 = mock.AsyncMock(spec=Node)
- self.server1 = GRPCServer(self.node1, config.peers["node1"].address, config.peers["node1"].port)
- self.server2 = GRPCServer(self.node2, config.peers["node2"].address, config.peers["node2"].port)
- await self.server1.start()
- await self.server2.start()
- self.discovery1 = ManualDiscovery(root_path, "node1", create_peer_handle=lambda peer_id, address, description, device_capabilities: GRPCPeerHandle(peer_id, address, description, device_capabilities))
- self.discovery2 = ManualDiscovery(root_path, "node2", create_peer_handle=lambda peer_id, address, description, device_capabilities: GRPCPeerHandle(peer_id, address, description, device_capabilities))
- await self.discovery1.start()
- await self.discovery2.start()
-
- async def asyncTearDown(self):
- await self.discovery1.stop()
- await self.discovery2.stop()
- await self.server1.stop()
- await self.server2.stop()
-
- async def test_grpc_discovery(self):
- peers1 = await self.discovery1.discover_peers(wait_for_peers=1)
- self.assertEqual(len(peers1), 1)
- peers2 = await self.discovery2.discover_peers(wait_for_peers=1)
- self.assertEqual(len(peers2), 1)
-
- # Connect
- await peers1[0].connect()
- await peers2[0].connect()
- self.assertTrue(await peers1[0].is_connected())
- self.assertTrue(await peers2[0].is_connected())
-
- # Kill server1
- await self.server1.stop()
-
- self.assertTrue(await peers1[0].is_connected())
- self.assertFalse(await peers2[0].is_connected())
-
- # Kill server2
- await self.server2.stop()
-
- self.assertFalse(await peers1[0].is_connected())
- self.assertFalse(await peers2[0].is_connected())
-
- async def test_dynamic_config_update(self):
- initial_peers = await self.discovery1.discover_peers(wait_for_peers=1)
- self.assertEqual(len(initial_peers), 1)
-
- # Save original config for cleanup
- with open(root_path, "r") as f:
- original_config = json.load(f)
-
- try:
- updated_config = {
- "peers": {
- **original_config["peers"],
- "node3": {
- "address": "localhost",
- "port": 50053,
- "device_capabilities": {"model": "Unknown Model", "chip": "Unknown Chip", "memory": 0, "flops": {"fp32": 0, "fp16": 0, "int8": 0}},
- },
- }
- }
-
- with open(root_path, "w") as f:
- json.dump(updated_config, f, indent=2)
-
- node3 = mock.AsyncMock(spec=Node)
- server3 = GRPCServer(node3, "localhost", 50053)
- await server3.start()
-
- try:
- # Wait for the config to be reloaded
+ async def asyncSetUp(self):
+ config = NetworkTopology.from_path(root_path)
+
+ self.node1 = mock.AsyncMock(spec=Node)
+ self.node2 = mock.AsyncMock(spec=Node)
+ self.server1 = GRPCServer(
+ self.node1, config.peers["node1"].address, config.peers["node1"].port
+ )
+ self.server2 = GRPCServer(
+ self.node2, config.peers["node2"].address, config.peers["node2"].port
+ )
+ await self.server1.start()
+ await self.server2.start()
+ self.discovery1 = ManualDiscovery(
+ root_path,
+ "node1",
+ create_peer_handle=lambda peer_id, address, description, device_capabilities: GRPCPeerHandle(
+ peer_id, address, description, device_capabilities
+ ),
+ )
+ self.discovery2 = ManualDiscovery(
+ root_path,
+ "node2",
+ create_peer_handle=lambda peer_id, address, description, device_capabilities: GRPCPeerHandle(
+ peer_id, address, description, device_capabilities
+ ),
+ )
+ await self.discovery1.start()
+ await self.discovery2.start()
+
+ async def asyncTearDown(self):
+ await self.discovery1.stop()
+ await self.discovery2.stop()
+ await self.server1.stop()
+ await self.server2.stop()
+
+ async def test_grpc_discovery(self):
+ peers1 = await self.discovery1.discover_peers(wait_for_peers=1)
+ self.assertEqual(len(peers1), 1)
+ peers2 = await self.discovery2.discover_peers(wait_for_peers=1)
+ self.assertEqual(len(peers2), 1)
+
+ # Connect
+ await peers1[0].connect()
+ await peers2[0].connect()
+ self.assertTrue(await peers1[0].is_connected())
+ self.assertTrue(await peers2[0].is_connected())
+
+ # Kill server1
+ await self.server1.stop()
+
+ self.assertTrue(await peers1[0].is_connected())
+ self.assertFalse(await peers2[0].is_connected())
+
+ # Kill server2
+ await self.server2.stop()
+
+ self.assertFalse(await peers1[0].is_connected())
+ self.assertFalse(await peers2[0].is_connected())
+
+ async def test_dynamic_config_update(self):
+ initial_peers = await self.discovery1.discover_peers(wait_for_peers=1)
+ self.assertEqual(len(initial_peers), 1)
+
+ # Save original config for cleanup
+ with open(root_path, "r") as f:
+ original_config = json.load(f)
+
+ try:
+ updated_config = {
+ "peers": {
+ **original_config["peers"],
+ "node3": {
+ "address": "localhost",
+ "port": 50053,
+ "device_capabilities": {
+ "model": "Unknown Model",
+ "chip": "Unknown Chip",
+ "memory": 0,
+ "flops": {"fp32": 0, "fp16": 0, "int8": 0},
+ },
+ },
+ }
+ }
+
+ with open(root_path, "w") as f:
+ json.dump(updated_config, f, indent=2)
+
+ node3 = mock.AsyncMock(spec=Node)
+ server3 = GRPCServer(node3, "localhost", 50053)
+ await server3.start()
+
+ try:
+ # Wait for the config to be reloaded
+ await asyncio.sleep(1.5)
+
+ updated_peers = await self.discovery1.discover_peers(wait_for_peers=2)
+ self.assertEqual(len(updated_peers), 2)
+
+ for peer in updated_peers:
+ await peer.connect()
+ self.assertTrue(await peer.is_connected())
+
+ finally:
+ await server3.stop()
+
+ finally:
+ # Restore the original config file
+ with open(root_path, "w") as f:
+ json.dump(original_config, f, indent=2)
+
+ # Wait for the config to be reloaded again
await asyncio.sleep(1.5)
- updated_peers = await self.discovery1.discover_peers(wait_for_peers=2)
- self.assertEqual(len(updated_peers), 2)
-
- for peer in updated_peers:
- await peer.connect()
- self.assertTrue(await peer.is_connected())
-
- finally:
- await server3.stop()
-
- finally:
- # Restore the original config file
- with open(root_path, "w") as f:
- json.dump(original_config, f, indent=2)
-
- # Wait for the config to be reloaded again
- await asyncio.sleep(1.5)
-
- updated_peers = await self.discovery1.discover_peers(wait_for_peers=1)
- self.assertEqual(len(updated_peers), 1)
-
+ updated_peers = await self.discovery1.discover_peers(wait_for_peers=1)
+ self.assertEqual(len(updated_peers), 1)
if __name__ == "__main__":
- asyncio.run(unittest.main())
+ asyncio.run(unittest.main())
← e5eb3259 handle when a peer is removed from config, so the known_peer
·
back to Exo
·
changes after rebase 90de7ead →