Skip to content

Commit 523ce17

Browse files
authored
Adding socket methods used by trio (#300)
* Adding `socket` methods used by `trio`, not 100% sure about the actual implementation. * Bump version.
1 parent e5b2275 commit 523ce17

File tree

4 files changed

+180
-1
lines changed

4 files changed

+180
-1
lines changed

mocket/__init__.py

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -31,4 +31,4 @@
3131
"FakeSSLContext",
3232
)
3333

34-
__version__ = "3.13.9"
34+
__version__ = "3.13.10"

mocket/socket.py

Lines changed: 82 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -191,6 +191,76 @@ def sendall(self, data, entry=None, *args, **kwargs):
191191
self.io.truncate()
192192
self.io.seek(0)
193193

194+
def sendmsg(
195+
self,
196+
buffers: list[ReadableBuffer],
197+
ancdata: list[tuple[int, bytes]] | None = None,
198+
flags: int = 0,
199+
address: Address | None = None,
200+
) -> int:
201+
if not buffers:
202+
return 0
203+
204+
data = b"".join(bytes(b) for b in buffers)
205+
self.sendall(data)
206+
return len(data)
207+
208+
def recvmsg(
209+
self,
210+
buffersize: int | None = None,
211+
ancbufsize: int | None = None,
212+
flags: int = 0,
213+
) -> tuple[bytes, list[tuple[int, bytes]]]:
214+
"""
215+
Receive a message from the socket.
216+
This is a mock implementation that reads from the MocketSocketIO.
217+
"""
218+
try:
219+
data = self.recv(buffersize)
220+
except BlockingIOError:
221+
return b"", []
222+
223+
# Mocking the ancillary data and flags as empty
224+
return data, []
225+
226+
def recvmsg_into(
227+
self,
228+
buffers: list[ReadableBuffer],
229+
ancbufsize: int | None = None,
230+
flags: int = 0,
231+
address: Address | None = None,
232+
):
233+
"""
234+
Receive a message into multiple buffers.
235+
This is a mock implementation that reads from the MocketSocketIO.
236+
"""
237+
if not buffers:
238+
return 0
239+
240+
try:
241+
data = self.recv(len(buffers[0]))
242+
except BlockingIOError:
243+
return 0
244+
245+
for i, buffer in enumerate(buffers):
246+
if i < len(data):
247+
buffer[: len(data)] = data
248+
else:
249+
buffer[:] = b""
250+
return len(data)
251+
252+
def recvfrom_into(
253+
self,
254+
buffer: WriteableBuffer,
255+
buffersize: int | None = None,
256+
flags: int | None = None,
257+
):
258+
"""
259+
Receive data into a buffer and return the number of bytes received.
260+
This is a mock implementation that reads from the MocketSocketIO.
261+
"""
262+
return self.recv_into(buffer, buffersize, flags), self._address
263+
194264
def recv_into(
195265
self,
196266
buffer: WriteableBuffer,
@@ -286,6 +356,18 @@ def send(
286356
self._entry = entry
287357
return len(data)
288358

359+
def accept(self) -> tuple[MocketSocket, _RetAddress]:
360+
"""Accept a connection and return a new MocketSocket object."""
361+
new_socket = MocketSocket(
362+
family=self._family,
363+
type=self._type,
364+
proto=self._proto,
365+
)
366+
new_socket._address = (self._host, self._port)
367+
new_socket._host = self._host
368+
new_socket._port = self._port
369+
return new_socket, (self._host, self._port)
370+
289371
def close(self) -> None:
290372
if self._true_socket and not self._true_socket._closed:
291373
self._true_socket.close()

pyproject.toml

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -61,6 +61,7 @@ test = [
6161
"mypy",
6262
"types-decorator",
6363
"types-requests",
64+
"trio",
6465
]
6566
speedups = [
6667
"xxhash;platform_python_implementation=='CPython'",

tests/test_socket.py

Lines changed: 96 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -30,3 +30,99 @@ def test_udp_socket():
3030

3131
assert data == response_data
3232
assert address == (host, port)
33+
34+
35+
def test_recvmsg():
36+
sock = MocketSocket(socket.AF_INET, socket.SOCK_STREAM)
37+
test_data = b"hello world"
38+
sock._io = type("MockIO", (), {"read": lambda self, n: test_data})()
39+
data, ancdata = sock.recvmsg(1024)
40+
assert data == test_data
41+
assert ancdata == []
42+
43+
44+
def test_recvmsg_into():
45+
sock = MocketSocket(socket.AF_INET, socket.SOCK_STREAM)
46+
test_data = b"foobar"
47+
sock._io = type("MockIO", (), {"read": lambda self, n: test_data})()
48+
buf = bytearray(10)
49+
buf2 = bytearray(10)
50+
buffers = [buf, buf2]
51+
nbytes = sock.recvmsg_into(buffers)
52+
assert nbytes == len(test_data)
53+
assert buf[: len(test_data)] == test_data
54+
55+
56+
def test_recvmsg_into_empty_buffers():
57+
sock = MocketSocket(socket.AF_INET, socket.SOCK_STREAM)
58+
result = sock.recvmsg_into([])
59+
assert result == 0
60+
61+
62+
def test_accept():
63+
sock = MocketSocket(socket.AF_INET, socket.SOCK_STREAM)
64+
sock._host = "127.0.0.1"
65+
sock._port = 8080
66+
new_sock, addr = sock.accept()
67+
assert isinstance(new_sock, MocketSocket)
68+
assert new_sock is not sock
69+
assert addr == ("127.0.0.1", 8080)
70+
assert new_sock._host == "127.0.0.1"
71+
assert new_sock._port == 8080
72+
73+
74+
@mocketize
75+
def test_sendmsg():
76+
sock = MocketSocket(socket.AF_INET, socket.SOCK_STREAM)
77+
sock._host = "127.0.0.1"
78+
sock._port = 8080
79+
response_data = b"pong"
80+
81+
Mocket.register(MocketEntry((sock._host, sock._port), [response_data]))
82+
83+
msg = [b"foo", b"bar", b"foobaz"]
84+
total_sent = sock.sendmsg(msg)
85+
assert total_sent == sum(len(m) for m in msg)
86+
assert Mocket.last_request() == b"".join(msg)
87+
88+
89+
def test_sendmsg_empty_buffers():
90+
sock = MocketSocket(socket.AF_INET, socket.SOCK_STREAM)
91+
result = sock.sendmsg([])
92+
assert result == 0
93+
94+
95+
def test_recvmsg_no_data():
96+
sock = MocketSocket(socket.AF_INET, socket.SOCK_STREAM)
97+
# Mock _io.read to return empty bytes
98+
sock._io = type("MockIO", (), {"read": lambda self, n: b""})()
99+
data, ancdata = sock.recvmsg(1024)
100+
assert data == b""
101+
assert ancdata == []
102+
103+
104+
def test_recvmsg_into_no_data():
105+
sock = MocketSocket(socket.AF_INET, socket.SOCK_STREAM)
106+
# Mock _io.read to return empty bytes
107+
sock._io = type("MockIO", (), {"read": lambda self, n: b""})()
108+
buf = bytearray(10)
109+
nbytes = sock.recvmsg_into([buf])
110+
assert nbytes == 0
111+
assert buf == bytearray(10)
112+
113+
114+
def test_getsockopt():
115+
# getsockopt is a static method, so we can call it directly
116+
result = MocketSocket.getsockopt(0, 0)
117+
assert result == socket.SOCK_STREAM
118+
119+
120+
def test_recvfrom_into():
121+
sock = MocketSocket(socket.AF_INET, socket.SOCK_STREAM)
122+
test_data = b"abc123"
123+
sock._io = type("MockIO", (), {"read": lambda self, n: test_data})()
124+
buf = bytearray(10)
125+
nbytes, addr = sock.recvfrom_into(buf)
126+
assert nbytes == len(test_data)
127+
assert buf[:nbytes] == test_data
128+
assert addr == sock._address

0 commit comments

Comments
 (0)