From e228f08fb47d3474037a77251c152dd5c1a40f8a 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..e8f6fec 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: + 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()