diff --git a/wish/python/src/web_stream_ext.cc b/wish/python/src/web_stream_ext.cc index 17b0d65..c4f8e14 100644 --- a/wish/python/src/web_stream_ext.cc +++ b/wish/python/src/web_stream_ext.cc @@ -385,11 +385,7 @@ NB_MODULE(web_stream_ext, m) { handler->SetOnMessage([&self](uint8_t opcode, const std::string& msg) { nb::gil_scoped_acquire acquire; try { - if (opcode == WEB_STREAM_OPCODE_BINARY || opcode == WEB_STREAM_OPCODE_METADATA) { - self.on_message_cb(opcode, nb::bytes(msg.data(), msg.size())); - } else { - self.on_message_cb(opcode, msg); - } + self.on_message_cb(opcode, nb::bytes(msg.data(), msg.size())); } catch (nb::python_error& e) { e.restore(); PyErr_WriteUnraisable(self.on_message_cb.ptr()); @@ -501,11 +497,7 @@ NB_MODULE(web_stream_ext, m) { handler->SetOnMessage([&self](uint8_t opcode, const std::string& msg) { nb::gil_scoped_acquire acquire; try { - if (opcode == WEB_STREAM_OPCODE_BINARY || opcode == WEB_STREAM_OPCODE_METADATA) { - self.on_message_cb(opcode, nb::bytes(msg.data(), msg.size())); - } else { - self.on_message_cb(opcode, msg); - } + self.on_message_cb(opcode, nb::bytes(msg.data(), msg.size())); } catch (nb::python_error& e) { e.restore(); PyErr_WriteUnraisable(self.on_message_cb.ptr()); diff --git a/wish/python/tests/test_e2e.py b/wish/python/tests/test_e2e.py index 8d6341a..628ac67 100644 --- a/wish/python/tests/test_e2e.py +++ b/wish/python/tests/test_e2e.py @@ -129,6 +129,123 @@ async def run(): asyncio.set_event_loop(loop) loop.run_until_complete(run()) + def test_recv_decode_options(self): + async def run(): + # Allow the server more time to start up to avoid Connection reset + await asyncio.sleep(1.0) + + uri = f"webstream://127.0.0.1:{self.port}" + + async with web_stream.connect(uri) as ws: + # 1. Test sending string and receiving with decode=False (returns bytes) + await ws.send("hello") + msg_bytes = await asyncio.wait_for(ws.recv(decode=False), timeout=2.0) + self.assertEqual(msg_bytes, b"hello") + self.assertIsInstance(msg_bytes, bytes) + + # 2. Test sending string and receiving with decode=True (returns str) + await ws.send("world") + msg_str = await asyncio.wait_for(ws.recv(decode=True), timeout=2.0) + self.assertEqual(msg_str, "world") + self.assertIsInstance(msg_str, str) + + # 3. Test sending bytes (which will be echoed as binary) and receiving with decode=True (returns str) + await ws.send(b"binary-text") + msg_str_from_bin = await asyncio.wait_for(ws.recv(decode=True), timeout=2.0) + self.assertEqual(msg_str_from_bin, "binary-text") + self.assertIsInstance(msg_str_from_bin, str) + + # 4. Test sending bytes and receiving with default decode=None (returns bytes for binary) + await ws.send(b"binary-bytes") + msg_bytes_from_bin = await asyncio.wait_for(ws.recv(decode=None), timeout=2.0) + self.assertEqual(msg_bytes_from_bin, b"binary-bytes") + self.assertIsInstance(msg_bytes_from_bin, bytes) + + loop = asyncio.new_event_loop() + asyncio.set_event_loop(loop) + loop.run_until_complete(run()) + + def test_recv_metadata(self): + async def run(): + # Allow the server more time to start up to avoid Connection reset + await asyncio.sleep(1.0) + + uri = f"webstream://127.0.0.1:{self.port}" + + async with web_stream.connect(uri) as ws: + # We can manually put a metadata message into ws._recv_queue + await ws._recv_queue.put((web_stream.WEB_STREAM_OPCODE_METADATA, b"meta-data-val")) + + # Default decode=None: should return bytes + msg = await ws.recv(decode=None) + self.assertEqual(msg, b"meta-data-val") + self.assertIsInstance(msg, bytes) + + # Put another one and test decode=True: should return str + await ws._recv_queue.put((web_stream.WEB_STREAM_OPCODE_METADATA, b"meta-data-val2")) + msg_str = await ws.recv(decode=True) + self.assertEqual(msg_str, "meta-data-val2") + self.assertIsInstance(msg_str, str) + + loop = asyncio.new_event_loop() + asyncio.set_event_loop(loop) + loop.run_until_complete(run()) + + def test_iterator_break(self): + async def run(): + # Allow the server more time to start up to avoid Connection reset + await asyncio.sleep(1.0) + + uri = f"webstream://127.0.0.1:{self.port}" + + async with web_stream.connect(uri) as ws: + await ws.send("hello") + await ws.send("world") + + received = [] + async for msg in ws: + received.append(msg) + if len(received) == 2: + break + + self.assertEqual(received, ["hello", "world"]) + + loop = asyncio.new_event_loop() + asyncio.set_event_loop(loop) + try: + loop.run_until_complete(run()) + finally: + loop.run_until_complete(loop.shutdown_asyncgens()) + loop.close() + + def test_iterator_close(self): + async def run(): + # Allow the server more time to start up to avoid Connection reset + await asyncio.sleep(1.0) + + uri = f"webstream://127.0.0.1:{self.port}" + + async with web_stream.connect(uri) as ws: + async def close_later(): + await asyncio.sleep(0.5) + await ws.close() + + asyncio.create_task(close_later()) + + messages = [] + async for msg in ws: + messages.append(msg) + + self.assertEqual(messages, []) + + loop = asyncio.new_event_loop() + asyncio.set_event_loop(loop) + try: + loop.run_until_complete(run()) + finally: + loop.run_until_complete(loop.shutdown_asyncgens()) + loop.close() + def test_server_disconnect_detect(self): port = get_free_port() cmd = [ diff --git a/wish/python/web_stream/__init__.py b/wish/python/web_stream/__init__.py index 4d5a628..79f4136 100644 --- a/wish/python/web_stream/__init__.py +++ b/wish/python/web_stream/__init__.py @@ -1,3 +1,15 @@ -from .client import connect, WebStreamConnection +from .client import ( + connect, + WebStreamConnection, + WEB_STREAM_OPCODE_TEXT, + WEB_STREAM_OPCODE_BINARY, + WEB_STREAM_OPCODE_METADATA, +) -__all__ = ["connect", "WebStreamConnection"] +__all__ = [ + "connect", + "WebStreamConnection", + "WEB_STREAM_OPCODE_TEXT", + "WEB_STREAM_OPCODE_BINARY", + "WEB_STREAM_OPCODE_METADATA", +] diff --git a/wish/python/web_stream/client.py b/wish/python/web_stream/client.py index ff9675c..f17bc28 100644 --- a/wish/python/web_stream/client.py +++ b/wish/python/web_stream/client.py @@ -4,6 +4,11 @@ from . import web_stream_ext +# WebStream Opcodes +WEB_STREAM_OPCODE_TEXT = 1 +WEB_STREAM_OPCODE_BINARY = 2 +WEB_STREAM_OPCODE_METADATA = 3 + class WebStreamConnection: def __init__(self, host, port, tls, ca_file="", cert_file="", key_file=""): self._host = host @@ -74,7 +79,7 @@ async def send(self, data): else: self._handler.send_text(str(data)) - async def recv(self): + async def recv(self, decode=None): """Receives a message from the WebStream connection.""" if not self._client: raise RuntimeError("Connection is closed") @@ -85,14 +90,32 @@ async def recv(self): raise res opcode, msg = res - # You can process opcode here if you want to distinguish text/binary - # We'll just return the message - # In actual implementation: 1=Text, 2=Binary - if opcode == 2: - return msg.encode("utf-8") if isinstance(msg, str) else msg + + # Determine if we should return str or bytes + # Default behavior: + # - WEB_STREAM_OPCODE_TEXT -> str + # - WEB_STREAM_OPCODE_BINARY -> bytes + # - WEB_STREAM_OPCODE_METADATA -> bytes + # Note: UTF-8 decoding for metadata is undefined in the specification. + # We treat it similarly to binary messages for now (i.e. no UTF-8 decoding by default). + should_decode = decode + if should_decode is None: + should_decode = (opcode == WEB_STREAM_OPCODE_TEXT) + + if should_decode: + return msg.decode("utf-8") else: return msg + async def __aiter__(self): + """Iterate on incoming messages.""" + try: + while True: + yield await self.recv() + except (ConnectionAbortedError, RuntimeError): + pass + + class _ConnectContextManager: def __init__(self, uri, ca_file="", cert_file="", key_file=""): parsed = urlparse(uri)