Raise exception on socket disconnect

Gr1N · Sep 5, 2019 · b43cecc · b43cecc
1 parent 0bfe172
commit b43cecc
Show file tree

Hide file tree

Showing 2 changed files with 15 additions and 0 deletions.
diff --git a/pynats/client.py b/pynats/client.py
@@ -272,6 +272,9 @@ def _readline(self, *, size: int = None) -> bytes:
             line = cast(bytes, self._socket_file.readline())
             read.write(line)
 
+            if len(line) == 0:
+                raise ConnectionResetError(self)
+
             if size is not None:
                 if read.tell() == size + len(_CRLF_):
                     break

diff --git a/tests/test_client.py b/tests/test_client.py
@@ -179,3 +179,15 @@ def test_request_timeout(nats_url):
     with NATSClient(nats_url, socket_timeout=2) as client:
         with pytest.raises(socket.timeout):
             client.request("test-subject")
+
+
+def test_exception_on_disconnect(nats_url):
+    with NATSClient(nats_url, socket_timeout=2) as client:
+        sub = client.subscribe(
+            "test-subject", callback=lambda x: x, queue="test-queue", max_messages=2
+        )
+
+        client._socket_file.readline = lambda: b""
+
+        with pytest.raises(ConnectionResetError):
+            client.wait()