Refactor radar system

locator_test_client.py

@@ -3,8 +3,8 @@ import random
 import socket
 import struct
 import threading
-import time
-from typing import Any, Dict, Tuple
+from contextlib import suppress
+from typing import Any
 
 HOST = "127.0.0.1"
 PORT = 8888
@@ -18,7 +18,7 @@ MAX_PAYLOAD_BYTES = 64 * 1024
 HEADER_STRUCT = struct.Struct("<II")
 
 
-def encode_packet(payload: Dict[str, Any], device_id: int) -> bytes:
+def encode_packet(payload: dict[str, Any], device_id: int) -> bytes:
     """Encode a JSON payload using the protocol binary header."""
     payload_bytes = json.dumps(
         payload,
@@ -28,12 +28,17 @@ def encode_packet(payload: Dict[str, Any], device_id: int) -> bytes:
     return HEADER_STRUCT.pack(device_id, len(payload_bytes)) + payload_bytes
 
 
-def recv_exactly(sock: socket.socket, size: int) -> bytes:
+def recv_exactly(sock: socket.socket, size: int, stop_event: threading.Event | None = None) -> bytes:
     """Receive an exact number of bytes from the socket."""
     chunks = bytearray()
 
     while len(chunks) < size:
-        chunk = sock.recv(size - len(chunks))
+        if stop_event is not None and stop_event.is_set():
+            raise ConnectionError("Client stopped.")
+        try:
+            chunk = sock.recv(size - len(chunks))
+        except socket.timeout:
+            continue  # Preserve a partial header/payload across timeout polls.
         if not chunk:
             raise ConnectionError("Connection closed by peer.")
         chunks.extend(chunk)
@@ -41,9 +46,9 @@ def recv_exactly(sock: socket.socket, size: int) -> bytes:
     return bytes(chunks)
 
 
-def read_packet(sock: socket.socket) -> Tuple[int, Any]:
+def read_packet(sock: socket.socket, stop_event: threading.Event | None = None) -> tuple[int, Any]:
     """Read and decode a single packet from the server."""
-    header_bytes = recv_exactly(sock, HEADER_STRUCT.size)
+    header_bytes = recv_exactly(sock, HEADER_STRUCT.size, stop_event)
     device_id, payload_length = HEADER_STRUCT.unpack(header_bytes)
 
     if payload_length > MAX_PAYLOAD_BYTES:
@@ -52,10 +57,11 @@ def read_packet(sock: socket.socket) -> Tuple[int, Any]:
             % (payload_length, MAX_PAYLOAD_BYTES)
         )
 
-    payload_bytes = recv_exactly(sock, payload_length)
+    payload_bytes = recv_exactly(sock, payload_length, stop_event)
     payload = json.loads(payload_bytes.decode("utf-8"))
     return device_id, payload
 
+
 def send_vlc_message(sock: socket.socket, device_id: int, vlc: float) -> None:
     """Send one test client payload to the server."""
     payload = {"vlc": vlc}
@@ -67,10 +73,8 @@ def receive_loop(sock: socket.socket, stop_event: threading.Event) -> None:
     """Print all packets received from the locator server."""
     while not stop_event.is_set():
         try:
-            device_id, payload = read_packet(sock)
-        except socket.timeout:
-            continue
-        except (ConnectionError, OSError, ValueError, json.JSONDecodeError) as error:
+            device_id, payload = read_packet(sock, stop_event)
+        except (OSError, ValueError) as error:
             if not stop_event.is_set():
                 print("Receiver stopped:", error)
             stop_event.set()
@@ -89,12 +93,10 @@ def send_loop(sock: socket.socket, stop_event: threading.Event) -> None:
     )
 
     if SEND_INTERVAL_SECONDS <= 0:
-        while not stop_event.is_set():
-            time.sleep(0.1)
+        stop_event.wait()
         return
 
-    while not stop_event.is_set():
-        time.sleep(SEND_INTERVAL_SECONDS)
+    while not stop_event.wait(SEND_INTERVAL_SECONDS):
         send_vlc_message(
             sock,
             CLIENT_DEVICE_ID,
@@ -124,8 +126,13 @@ def main() -> None:
             send_loop(sock, stop_event)
         except KeyboardInterrupt:
             print("Stopping client.")
+        except OSError as error:
+            print("Sender stopped:", error)
         finally:
             stop_event.set()
+            with suppress(OSError):
+                sock.shutdown(socket.SHUT_RDWR)
+            receiver.join(RECV_TIMEOUT_SECONDS + 1.0)
 
 
 if __name__ == "__main__":