Speed up test_autorelay

learning-at-home · Jun 11, 2023 · 654f5c4 · 654f5c4
1 parent 0d2614d
commit 654f5c4
Showing 1 changed file with 5 additions and 12 deletions.
diff --git a/tests/test_relays.py b/tests/test_relays.py
@@ -6,8 +6,8 @@
 import hivemind
 
 
-async def ping_to_client(dht, node, peer_id: str):
-    return await node.protocol.call_ping(hivemind.PeerID.from_base58(str(peer_id)))
+async def ping_to_client(dht, node, peer_id: hivemind.p2p.PeerID):
+    return await node.protocol.call_ping(peer_id)
 
 
 @pytest.mark.forked
@@ -38,7 +38,6 @@ def test_autorelay(use_auto_relay: bool, use_relay: bool):
         client_mode=False,
         use_auto_relay=use_auto_relay,
     )
-    time.sleep(5)
     dht_second_peer = hivemind.DHT(
         initial_peers=initial_peers,
         start=True,
@@ -50,15 +49,9 @@ def test_autorelay(use_auto_relay: bool, use_relay: bool):
 
     assert dht_first_peer.is_alive() and dht_second_peer.is_alive() and dht_third_peer.is_alive()
 
-    time_start = time.perf_counter()
-    while time.perf_counter() - time_start < 30:
-        reached_ip = dht_second_peer.run_coroutine(partial(ping_to_client, peer_id=dht_third_peer.peer_id))
-        if reached_ip:
-            assert use_relay
-            break
-        time.sleep(2)
-    else:
-        assert not use_relay
+    reached_ip = dht_second_peer.run_coroutine(partial(ping_to_client, peer_id=dht_third_peer.peer_id))
+    if reached_ip:
+        assert use_relay
 
     for peer in dht_first_peer, dht_second_peer, dht_third_peer:
         peer.shutdown()