Fix discovered server callback not being awaited (#660)

Author: caspervonb

nats/aio/client.py

@@ -1856,7 +1856,7 @@ def _process_disconnect(self) -> None:
         """
         self._status = Client.DISCONNECTED
 
-    def _process_info(
+    async def _process_info(
         self, info: Dict[str, Any], initial_connection: bool = False
     ) -> None:
         """
@@ -1899,7 +1899,7 @@ def _process_info(
 
                 if (not initial_connection and connect_urls
                         and self._discovered_server_cb):
-                    self._discovered_server_cb()
+                    await self._discovered_server_cb()
 
     def _host_is_ip(self, connect_url: Optional[str]) -> bool:
         if connect_url is None:
@@ -1960,7 +1960,7 @@ async def _process_connect_init(self) -> None:
         if srv_info.get("auth_required", False):
             self._auth_configured = True
 
-        self._process_info(srv_info, initial_connection=True)
+        await self._process_info(srv_info, initial_connection=True)
 
         if "version" in self._server_info:
             self._current_server.server_version = self._server_info["version"]

nats/protocol/parser.py

@@ -161,7 +161,7 @@ async def parse(self, data: bytes = b""):
                 if info:
                     info_line = info.groups()[0]
                     srv_info = json.loads(info_line.decode())
-                    self.nc._process_info(srv_info)
+                    await self.nc._process_info(srv_info)
                     del self.buf[:info.end()]
                     continue
 

tests/test_client.py

@@ -1922,13 +1922,8 @@ async def test_discover_servers_on_first_connect(self):
         await asyncio.sleep(1)
 
         options = {"servers": ["nats://127.0.0.1:4223", ]}
-
-        discovered_server_cb = mock.Mock()
-
-        with mock.patch("asyncio.iscoroutinefunction", return_value=True):
-            await nc.connect(
-                **options, discovered_server_cb=discovered_server_cb
-            )
+        discovered_server_cb = mock.AsyncMock()
+        await nc.connect(**options, discovered_server_cb=discovered_server_cb)
         self.assertTrue(nc.is_connected)
         await nc.close()
         self.assertTrue(nc.is_closed)
@@ -1941,11 +1936,8 @@ async def test_discover_servers_after_first_connect(self):
         nc = NATS()
 
         options = {"servers": ["nats://127.0.0.1:4223", ]}
-        discovered_server_cb = mock.Mock()
-        with mock.patch("asyncio.iscoroutinefunction", return_value=True):
-            await nc.connect(
-                **options, discovered_server_cb=discovered_server_cb
-            )
+        discovered_server_cb = mock.AsyncMock()
+        await nc.connect(**options, discovered_server_cb=discovered_server_cb)
 
         # Start rest of cluster members so that we receive them
         # connect_urls on the first connect.

tests/test_parser.py

@@ -33,7 +33,7 @@ async def _process_msg(self, sid, subject, reply, payload, headers=None):
     async def _process_err(self, err=None):
         pass
 
-    def _process_info(self, info):
+    async def _process_info(self, info):
         self._server_info = info