From b43cecca726b365a75be6626759804397c6d7db4 Mon Sep 17 00:00:00 2001 From: Vitali Falileev Date: Thu, 5 Sep 2019 13:11:29 +0300 Subject: [PATCH] Raise exception on socket disconnect --- pynats/client.py | 3 +++ tests/test_client.py | 12 ++++++++++++ 2 files changed, 15 insertions(+) diff --git a/pynats/client.py b/pynats/client.py index 8083063..97e0c25 100644 --- 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 index d4fadbe..c94cdcd 100644 --- 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()