fix: Use the ping() result, not only the timeout signal

lablup · Feb 28, 2024 · 488c479 · 488c479
1 parent bfc4805
commit 488c479
Showing 1 changed file with 5 additions and 7 deletions.
diff --git a/src/callosum/lower/zeromq.py b/src/callosum/lower/zeromq.py
@@ -413,13 +413,11 @@ async def __aenter__(self) -> ZeroMQRPCConnection:
         client_sock.connect(self.addr.uri)
         self._main_sock = client_sock
         self.transport._sock = client_sock
-        try:
-            await self.ping(
-                ping_timeout=int(self._handshake_timeout * 1000)
-                if self._handshake_timeout is not None
-                else 5000
-            )
-        except asyncio.TimeoutError:
+        if not await self.ping(
+            ping_timeout=int(self._handshake_timeout * 1000)
+            if self._handshake_timeout is not None
+            else 5000
+        ):
             raise AuthenticationError
         handshake_done = time.perf_counter()
         log.debug(