Skip to content

Commit 36da767

Browse files
committed
Add safe_close function for socket cleanup in tests
1 parent 0c8d2d0 commit 36da767

1 file changed

Lines changed: 25 additions & 10 deletions

File tree

tests/test_trustme.py

Lines changed: 25 additions & 10 deletions
Original file line numberDiff line numberDiff line change
@@ -284,6 +284,19 @@ def test_ca_from_pem(tmp_path: Path) -> None:
284284
assert ca1.private_key_pem.bytes() == ca2.private_key_pem.bytes()
285285

286286

287+
def safe_close(sock):
288+
if sock is None:
289+
return
290+
try:
291+
sock.shutdown(socket.SHUT_RDWR)
292+
except Exception:
293+
pass
294+
try:
295+
sock.close()
296+
except Exception:
297+
pass
298+
299+
287300
def check_connection_end_to_end(
288301
wrap_client: Callable[[CA, socket.socket, str], SslSocket],
289302
wrap_server: Callable[[LeafCert, socket.socket], SslSocket],
@@ -296,12 +309,12 @@ def fake_ssl_client(ca: CA, raw_client_sock: socket.socket, hostname: str) -> No
296309
# Send and receive some data to prove the connection is good
297310
wrapped_client_sock.send(b"x")
298311
assert wrapped_client_sock.recv(1) == b"y"
299-
wrapped_client_sock.close()
312+
safe_close(wrapped_client_sock)
300313
except: # pragma: no cover
301314
sys.excepthook(*sys.exc_info())
302315
raise
303316
finally:
304-
raw_client_sock.close()
317+
safe_close(raw_client_sock)
305318

306319
# Server side
307320
def fake_ssl_server(server_cert: LeafCert, raw_server_sock: socket.socket) -> None:
@@ -310,23 +323,25 @@ def fake_ssl_server(server_cert: LeafCert, raw_server_sock: socket.socket) -> No
310323
# Prove that we're connected
311324
assert wrapped_server_sock.recv(1) == b"x"
312325
wrapped_server_sock.send(b"y")
313-
wrapped_server_sock.close()
326+
safe_close(wrapped_server_sock)
314327
except: # pragma: no cover
315328
sys.excepthook(*sys.exc_info())
316329
raise
317330
finally:
318-
raw_server_sock.close()
331+
safe_close(raw_server_sock)
319332

320333
def doit(ca: CA, hostname: str, server_cert: LeafCert) -> None:
321334
# socketpair and ssl don't work together on py2, because... reasons.
322335
# So we need to do this the hard way.
323336
listener = socket.socket()
324-
listener.bind(("127.0.0.1", 0))
325-
listener.listen(1)
326-
raw_client_sock = socket.socket()
327-
raw_client_sock.connect(listener.getsockname())
328-
raw_server_sock, _ = listener.accept()
329-
listener.close()
337+
try:
338+
listener.bind(("127.0.0.1", 0))
339+
listener.listen(1)
340+
raw_client_sock = socket.socket()
341+
raw_client_sock.connect(listener.getsockname())
342+
raw_server_sock, _ = listener.accept()
343+
finally:
344+
safe_close(listener)
330345
with ThreadPoolExecutor(2) as tpe:
331346
f1 = tpe.submit(fake_ssl_client, ca, raw_client_sock, hostname)
332347
f2 = tpe.submit(fake_ssl_server, server_cert, raw_server_sock)

0 commit comments

Comments
 (0)