Fix HTTPS httpx mocking in wrap_bio() plus regressions occurred

mocket/io.py

@@ -3,9 +3,6 @@
 from __future__ import annotations
 
 import io
-import os
-
-from mocket.mocket import Mocket
 
 
 class MocketSocketIO(io.BytesIO):
@@ -21,17 +18,12 @@ def __init__(self, address: tuple) -> None:
         super().__init__()
 
     def write(self, content: bytes) -> int:
-        """Write content to the buffer and the pipe if available.
+        """Write content to the in-memory buffer.
 
         Args:
             content: Bytes to write
 
         Returns:
             Number of bytes written
         """
-        super().write(content)
-
-        _, w_fd = Mocket.get_pair(self._address)
-        if w_fd:
-            os.write(w_fd, content)
-        return len(content)
+        return super().write(content)

mocket/mocket.py

@@ -22,7 +22,10 @@
 class Mocket:
     """Singleton class managing all mock socket operations and entries."""
 
+    _socket_ios: ClassVar[dict[Address, Any]] = {}
     _socket_pairs: ClassVar[dict[Address, tuple[int, int]]] = {}
+    _pipe_uses_data: ClassVar[dict[Address, bool]] = {}
+    _pending_readables: ClassVar[dict[Address, int]] = {}
     _address: ClassVar[Address | tuple[None, None]] = (None, None)
     _entries: ClassVar[dict[Address, list[MocketEntry]]] = collections.defaultdict(list)
     _requests: ClassVar[list] = []
@@ -90,6 +93,39 @@ def set_pair(cls, address: Address, pair: tuple[int, int]) -> None:
         """
         cls._socket_pairs[address] = pair
 
+    @classmethod
+    def get_io(cls, address: Address) -> Any:
+        """Get the shared socket I/O buffer for an address, if any."""
+        return cls._socket_ios.get(address)
+
+    @classmethod
+    def set_io(cls, address: Address, socket_io: Any) -> None:
+        """Store the shared socket I/O buffer for an address."""
+        cls._socket_ios[address] = socket_io
+
+    @classmethod
+    def pipe_uses_data(cls, address: Address) -> bool:
+        """Return whether the pipe for an address carries mirrored response bytes."""
+        return cls._pipe_uses_data.get(address, False)
+
+    @classmethod
+    def set_pipe_uses_data(cls, address: Address, uses_data: bool) -> None:
+        """Store whether the pipe for an address carries mirrored response bytes."""
+        cls._pipe_uses_data[address] = uses_data
+
+    @classmethod
+    def get_pending_readables(cls, address: Address) -> int:
+        """Get the number of readiness bytes currently queued for an address."""
+        return cls._pending_readables.get(address, 0)
+
+    @classmethod
+    def set_pending_readables(cls, address: Address, count: int) -> None:
+        """Store the number of readiness bytes queued for an address."""
+        if count > 0:
+            cls._pending_readables[address] = count
+        else:
+            cls._pending_readables.pop(address, None)
+
     @classmethod
     def register(cls, *entries: MocketEntry) -> None:
         """Register mock entries with Mocket.
@@ -132,10 +168,15 @@ def collect(cls, data: Any) -> None:
     @classmethod
     def reset(cls) -> None:
         """Reset all Mocket state and clean up file descriptors."""
+        for socket_io in cls._socket_ios.values():
+            socket_io.close()
         for r_fd, w_fd in cls._socket_pairs.values():
             os.close(r_fd)
             os.close(w_fd)
+        cls._socket_ios = {}
         cls._socket_pairs = {}
+        cls._pipe_uses_data = {}
+        cls._pending_readables = {}
         cls._entries = collections.defaultdict(list)
         cls._requests = []
         cls._record_storage = None

mocket/socket.py

@@ -200,8 +200,14 @@ def proto(self) -> int:
     @property
     def io(self) -> MocketSocketIO:
         """Get or create the socket I/O object."""
-        if self._io is None:
-            self._io = MocketSocketIO((self._host, self._port))
+        if self._io is None or getattr(self._io, "closed", False):
+            address = self._address_key()
+            self._io = Mocket.get_io(address)
+            if self._io is not None and getattr(self._io, "closed", False):
+                self._io = None
+            if self._io is None:
+                self._io = MocketSocketIO(address)
+                Mocket.set_io(address, self._io)
         return self._io
 
     def fileno(self) -> int:
@@ -214,9 +220,76 @@ def fileno(self) -> int:
         r_fd, _ = Mocket.get_pair(address)
         if not r_fd:
             r_fd, w_fd = os.pipe()
+            os.set_blocking(r_fd, False)
+            os.set_blocking(w_fd, False)
             Mocket.set_pair(address, (r_fd, w_fd))
+        if self._io is not None and self._buffered_bytes():
+            if Mocket.pipe_uses_data(address):
+                self._mirror_buffer_to_pipe()
+            else:
+                self._sync_readable_pipe()
         return r_fd
 
+    def _address_key(self) -> Address:
+        """Return the current socket address tuple."""
+        return self._host, self._port
+
+    def _buffered_bytes(self) -> int:
+        """Return the number of unread bytes buffered in the socket I/O."""
+        return len(self.io.getvalue()) - self.io.tell()
+
+    def _clear_readable_pipe(self) -> None:
+        """Drain any stale readiness bytes from the pipe for this socket."""
+        address = self._address_key()
+        r_fd, _ = Mocket.get_pair(address)
+        if not r_fd:
+            return
+
+        while True:
+            try:
+                if not os.read(r_fd, self._buflen):
+                    break
+            except BlockingIOError:
+                break
+
+        Mocket.set_pending_readables(address, 0)
+
+    def _mirror_buffer_to_pipe(self) -> None:
+        """Mirror unread response bytes into the pipe for small payloads."""
+        address = self._address_key()
+        _, w_fd = Mocket.get_pair(address)
+        if not w_fd:
+            return
+
+        unread = self.io.getvalue()[self.io.tell() :]
+        if unread:
+            os.write(w_fd, unread)
+
+    def _sync_readable_pipe(self) -> None:
+        """Keep the readiness pipe in sync with the unread buffer size.
+
+        The pipe is used only to wake selector-based async clients. Response bytes
+        remain in the in-memory buffer so large payloads do not block on OS pipe
+        capacity.
+        """
+        address = self._address_key()
+        _, w_fd = Mocket.get_pair(address)
+        if not w_fd:
+            return
+
+        pending = Mocket.get_pending_readables(address)
+        desired = min(self._buffered_bytes(), self._buflen)
+        if desired <= pending:
+            return
+
+        try:
+            written = os.write(w_fd, b"\0" * (desired - pending))
+        except BlockingIOError:
+            written = 0
+
+        if written:
+            Mocket.set_pending_readables(address, pending + written)
+
     def gettimeout(self) -> float | None:
         """Get the socket timeout.
 
@@ -375,10 +448,20 @@ def sendall(
             response = self.true_sendall(data, *args, **kwargs)
 
         if response is not None:
+            address = self._address_key()
+            # Ensure the address pipe exists before deciding whether to mirror
+            # response bytes or only publish readiness signals.
+            self.fileno()
             self.io.seek(0)
+            self._clear_readable_pipe()
             self.io.write(response)
             self.io.truncate()
             self.io.seek(0)
+            Mocket.set_pipe_uses_data(address, len(response) <= self._buflen)
+            if Mocket.pipe_uses_data(address):
+                self._mirror_buffer_to_pipe()
+            else:
+                self._sync_readable_pipe()
 
     def sendmsg(
         self,
@@ -537,11 +620,32 @@ def recv(self, buffersize: int, flags: int | None = None) -> bytes:
         Raises:
             BlockingIOError: If socket is non-blocking and no data available
         """
-        r_fd, _ = Mocket.get_pair((self._host, self._port))
-        if r_fd:
-            return os.read(r_fd, buffersize)
+        if buffersize is None:
+            buffersize = self._buflen
+
+        address = self._address_key()
+        r_fd, _ = Mocket.get_pair(address)
+        if r_fd and Mocket.pipe_uses_data(address):
+            try:
+                pipe_data = os.read(r_fd, buffersize)
+            except BlockingIOError:
+                pipe_data = b""
+            if pipe_data:
+                # Keep in-memory buffer position in sync with bytes drained from the pipe.
+                self.io.seek(self.io.tell() + len(pipe_data))
+                return pipe_data
+
+        pending = Mocket.get_pending_readables(address)
+        if r_fd and self._buffered_bytes() and pending == 0:
+            self._sync_readable_pipe()
+            pending = Mocket.get_pending_readables(address)
+
         data = self.io.read(buffersize)
         if data:
+            if r_fd and pending:
+                drained = os.read(r_fd, min(len(data), pending))
+                Mocket.set_pending_readables(address, pending - len(drained))
+                self._sync_readable_pipe()
             return data
         # used by Redis mock
         exc = BlockingIOError()

mocket/ssl/context.py

@@ -4,6 +4,7 @@
 
 from typing import Any
 
+from mocket.mocket import Mocket
 from mocket.socket import MocketSocket
 from mocket.ssl.socket import MocketSSLSocket
 
@@ -111,7 +112,23 @@ def wrap_bio(
             MocketSSLSocket instance
         """
         ssl_obj = MocketSSLSocket()
-        ssl_obj._host = server_hostname
+        if isinstance(server_hostname, bytes):
+            hostname = server_hostname.decode("utf-8", errors="replace")
+        else:
+            hostname = server_hostname
+
+        current_address = Mocket._address
+        if isinstance(current_address, tuple) and len(current_address) == 2:
+            current_host, current_port = current_address
+        else:
+            current_host, current_port = None, None
+
+        effective_host = hostname if hostname is not None else current_host
+        ssl_obj._host = effective_host
+        if current_port is not None:
+            ssl_obj._port = current_port
+        if effective_host is not None and current_port is not None:
+            ssl_obj._address = (effective_host, current_port)
         return ssl_obj
 
 

mocket/ssl/socket.py

@@ -27,6 +27,9 @@ def __init__(self, *args: Any, **kwargs: Any) -> None:
 
         self._did_handshake: bool = False
         self._sent_non_empty_bytes: bool = False
+        self._has_written: bool = False
+        self._ssl_pending: bytes = b""
+        self._ssl_pending_pos: int = 0
         self._original_socket: MocketSocket = self
 
     def read(self, buffersize: int | None = None) -> bytes:
@@ -38,13 +41,28 @@ def read(self, buffersize: int | None = None) -> bytes:
         Returns:
             Bytes read from the socket
 
-        Raises:
-            ssl.SSLWantReadError: If handshake not completed and no data
         """
-        rv = self.io.read(buffersize)
+        if self._ssl_pending_pos < len(self._ssl_pending):
+            if buffersize is None:
+                rv = self._ssl_pending[self._ssl_pending_pos :]
+                self._ssl_pending_pos = len(self._ssl_pending)
+            else:
+                end = self._ssl_pending_pos + buffersize
+                rv = self._ssl_pending[self._ssl_pending_pos : end]
+                self._ssl_pending_pos = min(end, len(self._ssl_pending))
+        else:
+            rv = b""
         if rv:
             self._sent_non_empty_bytes = True
-        if self._did_handshake and not self._sent_non_empty_bytes:
+
+        # asyncio SSL transports probe reads before writing request bytes.
+        # Keep that non-blocking behavior, but once writes happened we must
+        # return empty bytes instead of surfacing SSLWantReadError.
+        if (
+            self._did_handshake
+            and not self._sent_non_empty_bytes
+            and not self._has_written
+        ):
             raise ssl.SSLWantReadError("The operation did not complete (read)")
         return rv
 
@@ -57,7 +75,14 @@ def write(self, data: bytes) -> int | None:
         Returns:
             Number of bytes written
         """
-        return self.send(encode_to_bytes(data))
+        self._has_written = self._has_written or bool(data)
+        bytes_sent = self.send(encode_to_bytes(data))
+
+        # Keep a private read buffer for SSL protocol consumers so response
+        # parsing does not depend on shared socket I/O cursor state.
+        self._ssl_pending = self.io.getvalue()
+        self._ssl_pending_pos = 0
+        return bytes_sent
 
     def do_handshake(self) -> None:
         """Perform SSL handshake (mock implementation)."""
@@ -72,8 +97,13 @@ def getpeercert(self, binary_form: bool = False) -> _PeerCertRetDictType:
         Returns:
             Mock certificate dictionary
         """
-        if not (self._host and self._port):
-            self._address = self._host, self._port = Mocket._address
+        if self._host is None or self._port is None:
+            current_host, current_port = Mocket._address
+            if self._host is None:
+                self._host = current_host
+            if self._port is None:
+                self._port = current_port
+            self._address = (self._host, self._port)
 
         now = datetime.now()
         shift = now + timedelta(days=30 * 12)
@@ -156,5 +186,8 @@ def _create(
 
         ssl_socket._io = sock._io
         ssl_socket._entry = sock._entry
+        ssl_socket._has_written = getattr(sock, "_has_written", False)
+        ssl_socket._ssl_pending = getattr(sock, "_ssl_pending", b"")
+        ssl_socket._ssl_pending_pos = getattr(sock, "_ssl_pending_pos", 0)
 
         return ssl_socket