Improved receive messages logic

Gr1N · Gr1N · commit b0124add9c57 · 2018-09-21T13:51:59.000+03:00
diff --git a/CHANGELOG.md b/CHANGELOG.md
@@ -2,7 +2,7 @@
 
 ## 0.3.0 (2018-XX-XX)
 
-- Improved receive messages logic: `.strip()` replaced with `.replace()`
+- Improved receive messages logic
 - Fixed regex deprecation warnings
 
 ## 0.2.0 (2018-09-01)
diff --git a/pynats/client.py b/pynats/client.py
@@ -265,26 +265,35 @@ def _recv(self, *commands: Pattern[bytes]) -> Tuple[Pattern[bytes], Match[bytes]
 
         return command, result
 
-    def _readline(self) -> bytes:
-        lines = []
+    def _readline(self, *, size: int = None) -> bytes:
+        read = b""
 
         while True:
             line = cast(bytes, self._socket_file.readline())
-            lines.append(line)
+            read += line
 
-            if line.endswith(_CRLF_):
+            if size is not None:
+                if len(self._strip(read)) == size:
+                    break
+            elif line.endswith(_CRLF_):
                 break
 
-        return b"".join(lines)
+        return read
+
+    def _strip(self, line: bytes) -> bytes:
+        return line[: -len(_CRLF_)]
 
     def _get_command(self, line: bytes) -> Optional[Pattern[bytes]]:
-        values = line.replace(_CRLF_, b"").split(b" ", 1)
+        values = self._strip(line).split(b" ", 1)
 
         return COMMANDS.get(values[0])
 
     def _handle_message(self, result: Match[bytes]) -> None:
         message_data = result.groupdict()
-        message_payload = self._readline().replace(_CRLF_, b"")
+
+        message_payload_size = int(message_data["size"])
+        message_payload = self._readline(size=message_payload_size)
+        message_payload = self._strip(message_payload)
 
         message = NATSMessage(
             sid=int(message_data["sid"].decode()),
diff --git a/tests/test_client.py b/tests/test_client.py
@@ -136,7 +136,12 @@ def worker():
         with NATSClient(socket_timeout=2) as client:
 
             def callback(message):
-                client.publish(message.reply, payload=msgpack.packb(32))
+                client.publish(
+                    message.reply,
+                    payload=msgpack.packb(
+                        {b"v": 3338} if message.payload else {b"v": 32}
+                    ),
+                )
 
             client.subscribe(
                 "test-subject", callback=callback, queue="test-queue", max_messages=2
@@ -153,13 +158,13 @@ def callback(message):
         resp = client.request("test-subject")
         assert resp.subject.startswith("_INBOX.")
         assert resp.reply == ""
-        assert msgpack.unpackb(resp.payload) == 32
+        assert msgpack.unpackb(resp.payload) == {b"v": 32}
 
         # request with payload
         resp = client.request("test-subject", payload=msgpack.packb("test-payload"))
         assert resp.subject.startswith("_INBOX.")
         assert resp.reply == ""
-        assert msgpack.unpackb(resp.payload) == 32
+        assert msgpack.unpackb(resp.payload) == {b"v": 3338}
 
     t.join()
 
