From bd690d8861639cecb76174067dce5b0adf438d98 Mon Sep 17 00:00:00 2001
From: Shubham Padkonde <shubhampadkonde12@gmail.com>
Date: Tue, 29 Sep 2026 19:41:06 +0530
Subject: [PATCH 1/2] fix: handle partial WLAN transfers and closed connections

Co-authored-by: Codex <noreply@openai.com>
---
 pslab/connection/wlan.py | 19 +++++++++--
 tests/test_wlan.py       | 73 ++++++++++++++++++++++++++++++++++++++++
 2 files changed, 90 insertions(+), 2 deletions(-)
 create mode 100644 tests/test_wlan.py

diff --git a/pslab/connection/wlan.py b/pslab/connection/wlan.py
index 79cde7a..009c3e6 100644
--- a/pslab/connection/wlan.py
+++ b/pslab/connection/wlan.py
@@ -77,6 +77,11 @@ def read(self, numbytes: int) -> bytes:
         Returns
         -------
         data : bytes
+
+        Raises
+        ------
+        ConnectionError
+            If the connection closes before all requested bytes are received.
         """
         received = b""
         buf_size = 4096
@@ -84,6 +89,8 @@ def read(self, numbytes: int) -> bytes:
 
         while remaining > 0:
             chunk = self._sock.recv(min(remaining, buf_size))
+            if not chunk:
+                raise ConnectionError("Connection closed while reading from PSLab.")
             received += chunk
             remaining -= len(chunk)
 
@@ -100,6 +107,11 @@ def write(self, data: bytes) -> int:
         -------
         numbytes : int
             Number of bytes written.
+
+        Raises
+        ------
+        ConnectionError
+            If the connection closes before all bytes are sent.
         """
         buf_size = 4096
         remaining = len(data)
@@ -107,8 +119,11 @@ def write(self, data: bytes) -> int:
 
         while remaining > 0:
             chunk = data[sent : sent + min(remaining, buf_size)]
-            sent += self._sock.send(chunk)
-            remaining -= len(chunk)
+            count = self._sock.send(chunk)
+            if not count:
+                raise ConnectionError("Connection closed while writing to PSLab.")
+            sent += count
+            remaining -= count
 
         return sent
 
diff --git a/tests/test_wlan.py b/tests/test_wlan.py
new file mode 100644
index 0000000..83b872c
--- /dev/null
+++ b/tests/test_wlan.py
@@ -0,0 +1,73 @@
+"""Test wireless transfers without a connected PSLab."""
+
+import socket
+from unittest.mock import Mock
+
+import pytest
+
+from pslab.connection.wlan import WLANHandler
+
+
+@pytest.fixture
+def handler():
+    """Return a handler with a controlled socket."""
+    device = WLANHandler()
+    device.disconnect()
+    device._sock = Mock()
+    return device
+
+
+@pytest.mark.parametrize("size", [10, 5000])
+def test_write_partial_sends(handler, size):
+    """Send every byte even when the socket accepts only part of a chunk."""
+    received = bytearray()
+    data = bytes(index % 256 for index in range(size))
+
+    def send(chunk):
+        count = min(len(chunk), 3)
+        received.extend(chunk[:count])
+        return count
+
+    handler._sock.send.side_effect = send
+    assert handler.write(data) == len(data)
+    assert received == data
+
+
+@pytest.mark.parametrize("prefix", [b"", b"abc"])
+def test_read_closed_connection(handler, prefix):
+    """Stop at EOF instead of retrying a closed socket indefinitely."""
+    responses = [prefix] if prefix else []
+    handler._sock.recv.side_effect = responses + [
+        b"",
+        AssertionError("Read continued after EOF"),
+    ]
+    with pytest.raises(ConnectionError):
+        handler.read(10)
+
+
+def test_write_closed_connection(handler):
+    """Reject a zero-byte send instead of reporting a truncated transfer."""
+    handler._sock.send.return_value = 0
+    with pytest.raises(ConnectionError):
+        handler.write(b"abc")
+
+
+def test_socket_transfer(handler):
+    """Exercise both directions over real connected sockets."""
+    client, peer = socket.socketpair()
+    with client, peer:
+        client.settimeout(1)
+        peer.settimeout(1)
+        handler._sock = client
+        peer.sendall(b"response")
+        assert handler.read(8) == b"response"
+        assert handler.write(b"request") == 7
+        assert peer.recv(7) == b"request"
+
+
+def test_empty_transfer(handler):
+    """Zero-length transfers require no socket access."""
+    assert handler.read(0) == b""
+    assert handler.write(b"") == 0
+    handler._sock.recv.assert_not_called()
+    handler._sock.send.assert_not_called()

From 54cc39387e49a8c55ae804c9ca5697f114579304 Mon Sep 17 00:00:00 2001
From: Shubham Padkonde <shubhampadkonde12@gmail.com>
Date: Tue, 29 Sep 2026 23:20:33 +0530
Subject: [PATCH 2/2] test: read complete WLAN socket response

---
 tests/test_wlan.py | 3 ++-
 1 file changed, 2 insertions(+), 1 deletion(-)

diff --git a/tests/test_wlan.py b/tests/test_wlan.py
index 83b872c..b5b7527 100644
--- a/tests/test_wlan.py
+++ b/tests/test_wlan.py
@@ -62,7 +62,8 @@ def test_socket_transfer(handler):
         peer.sendall(b"response")
         assert handler.read(8) == b"response"
         assert handler.write(b"request") == 7
-        assert peer.recv(7) == b"request"
+        with peer.makefile("rb") as stream:
+            assert stream.read(7) == b"request"
 
 
 def test_empty_transfer(handler):