Refactor message parsing functions and update tests

- Remove parse_msg_args and parse_hmsg_args in favor of async parse_msg and
parse_hmsg - Rename and update protocol parsing functions for clarity - Update
tests to use new async parsing functions

Author: caspervonb

nats-client/src/nats/client/protocol/message.py

@@ -147,80 +147,7 @@ def parse_control_line(line: bytes) -> tuple[str, list[str]]:
         raise ParseError(msg) from e
 
 
-def parse_msg_args(args: list[str]) -> tuple[str, str, str | None, int]:
-    """Parse MSG arguments into components.
 
-    Args:
-        args: MSG command arguments
-
-    Returns:
-        Tuple of (subject, sid, reply_to, payload_size)
-
-    Raises:
-        ParseError: If arguments are invalid
-    """
-    match len(args):
-        case 0 | 1 | 2:
-            msg = "Invalid MSG: not enough arguments"
-            raise ParseError(msg)
-        case 3:
-            subject, sid, size_str = args
-            try:
-                size = int(size_str)
-            except ValueError as e:
-                msg = f"Invalid payload size: {size_str}"
-                raise ParseError(msg) from e
-            return subject, sid, None, size
-        case 4:
-            subject, sid, reply_to, size_str = args
-            try:
-                size = int(size_str)
-            except ValueError as e:
-                msg = f"Invalid payload size: {size_str}"
-                raise ParseError(msg) from e
-            return subject, sid, reply_to, size
-        case _:
-            msg = "Invalid MSG: too many arguments"
-            raise ParseError(msg)
-
-
-def parse_hmsg_args(args: list[str]) -> tuple[str, str, str, int, int]:
-    """Parse HMSG arguments into components.
-
-    Args:
-        args: HMSG command arguments
-
-    Returns:
-        Tuple of (subject, sid, reply_to, header_size, total_size)
-
-    Raises:
-        ParseError: If arguments are invalid
-    """
-    match len(args):
-        case 0 | 1 | 2 | 3 | 4:
-            msg = "Invalid HMSG: not enough arguments"
-            raise ParseError(msg)
-        case 5:
-            subject, sid, reply_to, header_size_str, total_size_str = args
-            try:
-                header_size = int(header_size_str)
-                total_size = int(total_size_str)
-            except ValueError as e:
-                msg = f"Invalid size values: {header_size_str}, {total_size_str}"
-                raise ParseError(msg) from e
-
-            if header_size > MAX_HEADER_SIZE:
-                msg = f"Header too large: {header_size} > {MAX_HEADER_SIZE}"
-                raise ParseError(msg)
-
-            if header_size > total_size:
-                msg = f"Header size {header_size} larger than total size {total_size}"
-                raise ParseError(msg)
-
-            return subject, sid, reply_to, header_size, total_size
-        case _:
-            msg = "Invalid HMSG: too many arguments"
-            raise ParseError(msg)
 
 
 def parse_headers(
@@ -320,11 +247,11 @@ def parse_err(text: str) -> str:
     return text
 
 
-async def _parse_msg(reader: Reader, args: list[bytes]) -> Msg:
+async def parse_msg(reader: Reader, args: list[bytes]) -> Msg:
     """Parse MSG message.
 
     Args:
-        reader: AsyncIO stream reader
+        reader: Reader protocol implementation
         args: Message arguments
 
     Returns:
@@ -367,11 +294,11 @@ async def _parse_msg(reader: Reader, args: list[bytes]) -> Msg:
     return Msg("MSG", subject, sid, reply_to, payload)
 
 
-async def _parse_hmsg(reader: Reader, args: list[bytes]) -> HMsg:
+async def parse_hmsg(reader: Reader, args: list[bytes]) -> HMsg:
     """Parse HMSG message.
 
     Args:
-        reader: AsyncIO stream reader
+        reader: Reader protocol implementation
         args: Message arguments
 
     Returns:
@@ -432,7 +359,7 @@ async def _parse_hmsg(reader: Reader, args: list[bytes]) -> HMsg:
     )
 
 
-async def _parse_info(args: list[bytes]) -> Info:
+async def parse_info(args: list[bytes]) -> Info:
     """Parse INFO message.
 
     Args:
@@ -460,7 +387,7 @@ async def _parse_info(args: list[bytes]) -> Info:
         raise ParseError(msg) from e
 
 
-async def _parse_err(args: list[bytes]) -> Err:
+async def parse_err(args: list[bytes]) -> Err:
     """Parse ERR message.
 
     Args:
@@ -487,20 +414,20 @@ async def _parse_err(args: list[bytes]) -> Err:
     return Err("ERR", error_text)
 
 
-async def _parse_ping() -> Ping:
-    """Parse PING message.
+async def ping() -> Ping:
+    """Create PING message.
 
     Returns:
-        Parsed PING message
+        PING message
     """
     return Ping("PING")
 
 
-async def _parse_pong() -> Pong:
-    """Parse PONG message.
+async def pong() -> Pong:
+    """Create PONG message.
 
     Returns:
-        Parsed PONG message
+        PONG message
     """
     return Pong("PONG")
 
@@ -509,7 +436,7 @@ async def parse(reader: Reader) -> Message | None:
     """Parse a message from the protocol stream.
 
     Args:
-        reader: AsyncIO stream reader
+        reader: Reader protocol implementation
 
     Returns:
         Parsed protocol message or None if connection closed
@@ -545,17 +472,17 @@ async def parse(reader: Reader) -> Message | None:
 
         match op:
             case b"MSG":
-                return await _parse_msg(reader, args)
+                return await parse_msg(reader, args)
             case b"HMSG":
-                return await _parse_hmsg(reader, args)
+                return await parse_hmsg(reader, args)
             case b"PING":
-                return await _parse_ping()
+                return await ping()
             case b"PONG":
-                return await _parse_pong()
+                return await pong()
             case b"INFO":
-                return await _parse_info(args)
+                return await parse_info(args)
             case b"ERR":
-                return await _parse_err(args)
+                return await parse_err(args)
             case _:
                 # Use repr for better error reporting with control characters
                 msg = f"Unknown operation: {op!r}"

nats-client/tests/test_protocol.py

@@ -1,5 +1,6 @@
 """Tests for NATS protocol message parsing and command encoding."""
 
+import asyncio
 import json
 
 import pytest
@@ -16,8 +17,8 @@
     ParseError,
     parse_control_line,
     parse_headers,
-    parse_hmsg_args,
-    parse_msg_args,
+    parse_hmsg,
+    parse_msg,
 )
 from nats.client.protocol.types import ConnectInfo
 
@@ -72,69 +73,85 @@ def test_parse_control_line():
         parse_control_line(b"MSG " + b"x" * 4096)
 
 
-def test_parse_msg_args():
-    """Test parsing MSG arguments."""
+@pytest.mark.asyncio
+async def test_parse_msg():
+    """Test parsing MSG messages."""
     # Test valid MSG without reply
-    subject, sid, reply_to, size = parse_msg_args(["foo.bar", "1", "42"])
-    assert subject == "foo.bar"
-    assert sid == "1"
-    assert reply_to is None
-    assert size == 42
+    reader = asyncio.StreamReader()
+    reader.feed_data(b"hello\r\n")
+    reader.feed_eof()
+
+    msg = await parse_msg(reader, [b"foo.bar", b"1", b"5"])
+    assert msg.subject == "foo.bar"
+    assert msg.sid == "1"
+    assert msg.reply_to is None
+    assert msg.payload == b"hello"
 
     # Test valid MSG with reply
-    subject, sid, reply_to, size = parse_msg_args([
-        "foo.bar", "1", "reply.to", "42"
-    ])
-    assert subject == "foo.bar"
-    assert sid == "1"
-    assert reply_to == "reply.to"
-    assert size == 42
+    reader = asyncio.StreamReader()
+    reader.feed_data(b"hello\r\n")
+    reader.feed_eof()
+
+    msg = await parse_msg(reader, [b"foo.bar", b"1", b"reply.to", b"5"])
+    assert msg.subject == "foo.bar"
+    assert msg.sid == "1"
+    assert msg.reply_to == "reply.to"
+    assert msg.payload == b"hello"
 
     # Test invalid size
-    with pytest.raises(ParseError, match="Invalid payload size"):
-        parse_msg_args(["foo.bar", "1", "invalid"])
+    reader = asyncio.StreamReader()
+    with pytest.raises(ValueError):
+        await parse_msg(reader, [b"foo.bar", b"1", b"invalid"])
 
     # Test not enough arguments
+    reader = asyncio.StreamReader()
     with pytest.raises(ParseError, match="Invalid MSG: not enough arguments"):
-        parse_msg_args(["foo.bar", "1"])
+        await parse_msg(reader, [b"foo.bar", b"1"])
 
-    # Test too many arguments
-    with pytest.raises(ParseError, match="Invalid MSG: too many arguments"):
-        parse_msg_args(["foo.bar", "1", "reply.to", "42", "extra"])
+    # Test payload too large
+    reader = asyncio.StreamReader()
+    with pytest.raises(ParseError, match="Payload too large"):
+        await parse_msg(reader, [b"foo.bar", b"1", b"67108865"])
 
 
-def test_parse_hmsg_args():
-    """Test parsing HMSG arguments."""
+@pytest.mark.asyncio
+async def test_parse_hmsg():
+    """Test parsing HMSG messages."""
     # Test valid HMSG
-    subject, sid, reply_to, header_size, total_size = parse_hmsg_args([
-        "foo.bar", "1", "reply.to", "10", "52"
-    ])
-    assert subject == "foo.bar"
-    assert sid == "1"
-    assert reply_to == "reply.to"
-    assert header_size == 10
-    assert total_size == 52
+    reader = asyncio.StreamReader()
+    header_data = b"NATS/1.0\r\n\r\n"
+    payload = b"hello"
+    reader.feed_data(header_data + payload + b"\r\n")
+    reader.feed_eof()
+
+    header_size = len(header_data)
+    total_size = header_size + len(payload)
+    msg = await parse_hmsg(reader, [b"foo.bar", b"1", b"reply.to", str(header_size).encode(), str(total_size).encode()])
+    assert msg.subject == "foo.bar"
+    assert msg.sid == "1"
+    assert msg.reply_to == "reply.to"
+    assert msg.payload == b"hello"
+    assert msg.headers == {}
 
     # Test invalid sizes
-    with pytest.raises(ParseError, match="Invalid size values"):
-        parse_hmsg_args(["foo.bar", "1", "reply.to", "invalid", "52"])
+    reader = asyncio.StreamReader()
+    with pytest.raises(ValueError):
+        await parse_hmsg(reader, [b"foo.bar", b"1", b"reply.to", b"invalid", b"52"])
 
     # Test header size too large
-    with pytest.raises(ParseError, match="Header too large"):
-        parse_hmsg_args(["foo.bar", "1", "reply.to", "65537", "65538"])
+    reader = asyncio.StreamReader()
+    with pytest.raises(ParseError, match="Headers too large"):
+        await parse_hmsg(reader, [b"foo.bar", b"1", b"reply.to", b"65537", b"65538"])
 
-    # Test header size larger than total
-    with pytest.raises(ParseError,
-                       match="Header size .* larger than total size"):
-        parse_hmsg_args(["foo.bar", "1", "reply.to", "52", "10"])
+    # Test total size too large
+    reader = asyncio.StreamReader()
+    with pytest.raises(ParseError, match="Total message too large"):
+        await parse_hmsg(reader, [b"foo.bar", b"1", b"reply.to", b"10", b"67108865"])
 
     # Test not enough arguments
+    reader = asyncio.StreamReader()
     with pytest.raises(ParseError, match="Invalid HMSG: not enough arguments"):
-        parse_hmsg_args(["foo.bar", "1", "reply.to", "10"])
-
-    # Test too many arguments
-    with pytest.raises(ParseError, match="Invalid HMSG: too many arguments"):
-        parse_hmsg_args(["foo.bar", "1", "reply.to", "10", "52", "extra"])
+        await parse_hmsg(reader, [b"foo.bar", b"1", b"10"])
 
 
 def test_parse_headers():