Support for UDP sockets

.pre-commit-config.yaml

@@ -15,7 +15,7 @@ repos:
         exclude: helm/
         args: [ --unsafe ]
   - repo: https://github.com/charliermarsh/ruff-pre-commit
-    rev: "v0.11.2"
+    rev: "v0.11.11"
     hooks:
       - id: ruff
         args: [--fix, --exit-non-zero-on-fix]

mocket/socket.py

@@ -170,6 +170,11 @@ def makefile(self, mode: str = "r", bufsize: int = -1) -> MocketSocketIO:
     def get_entry(self, data: bytes) -> MocketEntry | None:
         return Mocket.get_entry(self._host, self._port, data)
 
+    def sendto(self, data: ReadableBuffer, address: Address | None = None) -> int:
+        self.connect(address)
+        self.sendall(data)
+        return len(data)
+
     def sendall(self, data, entry=None, *args, **kwargs):
         if entry is None:
             entry = self.get_entry(data)
@@ -204,6 +209,11 @@ def recv_into(
             buffer[: len(data)] = data
         return len(data)
 
+    def recvfrom(
+        self, buffersize: int, flags: int | None = None
+    ) -> tuple[bytes, _RetAddress]:
+        return self.recv(buffersize, flags), self._address
+
     def recv(self, buffersize: int, flags: int | None = None) -> bytes:
         r_fd, _ = Mocket.get_pair((self._host, self._port))
         if r_fd:

tests/test_socket.py

@@ -2,6 +2,7 @@
 
 import pytest
 
+from mocket import Mocket, MocketEntry, mocketize
 from mocket.socket import MocketSocket
 
 
@@ -11,3 +12,21 @@ def test_blocking_socket(blocking):
     sock.connect(("locahost", 1234))
     sock.setblocking(blocking)
     assert sock.getblocking() is blocking
+
+
+@mocketize
+def test_udp_socket():
+    host = "127.0.0.1"
+    port = 9999
+    request_data = b"ping"
+    response_data = b"pong"
+
+    Mocket.register(MocketEntry((host, port), [response_data]))
+
+    # Your UDP client code
+    sock = socket.socket(socket.AF_INET, socket.SOCK_DGRAM)
+    sock.sendto(request_data, (host, port))
+    data, address = sock.recvfrom(1024)
+
+    assert data == response_data
+    assert address == (host, port)