Skip to content

Commit a77d768

Browse files
committed
test(psycopg): add unit tests for happy path and error cleanup paths to reach 100% coverage
1 parent a30f5ad commit a77d768

1 file changed

Lines changed: 309 additions & 0 deletions

File tree

‎tests/unit/test_psycopg.py‎

Lines changed: 309 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -422,3 +422,312 @@ def mock_connect_impl(*args: Any, **kwargs: Any) -> MagicMock:
422422
assert conn is not None
423423
assert mock_remove.called
424424
assert mock_rmdir.called
425+
426+
427+
def test_proxy_happy_path_sequential() -> None:
428+
"""Test that _proxy forwards sequential request/response traffic cleanly without errors."""
429+
local_client, local_proxy = socket.socketpair(
430+
socket.AF_UNIX, socket.SOCK_STREAM
431+
)
432+
remote_proxy, remote_server = socket.socketpair(
433+
socket.AF_UNIX, socket.SOCK_STREAM
434+
)
435+
436+
t = threading.Thread(
437+
target=_proxy, args=(local_proxy, remote_proxy), daemon=True
438+
)
439+
t.start()
440+
441+
# Step 1: Client sends query
442+
local_client.sendall(b"SELECT 1;")
443+
received_query = remote_server.recv(1024)
444+
assert received_query == b"SELECT 1;"
445+
446+
# Step 2: Server sends response
447+
remote_server.sendall(b"RESULT 1")
448+
received_resp = local_client.recv(1024)
449+
assert received_resp == b"RESULT 1"
450+
451+
# Step 3: Client closes connection cleanly
452+
local_client.shutdown(socket.SHUT_RDWR)
453+
local_client.close()
454+
remote_server.close()
455+
t.join(timeout=2.0)
456+
assert not t.is_alive()
457+
458+
459+
def test_proxy_backpressure_and_clean_teardown() -> None:
460+
"""Demonstrate that Cloud SQL's single-threaded selector proxy exits cleanly
461+
462+
without deadlocking on teardown.
463+
"""
464+
local_client, local_proxy = socket.socketpair(
465+
socket.AF_UNIX, socket.SOCK_STREAM
466+
)
467+
remote_proxy, remote_server = socket.socketpair(
468+
socket.AF_UNIX, socket.SOCK_STREAM
469+
)
470+
471+
t_proxy = threading.Thread(
472+
target=_proxy, args=(local_proxy, remote_proxy), daemon=True
473+
)
474+
t_proxy.start()
475+
476+
# Step 1: Client sends query
477+
local_client.sendall(b"SELECT 1;")
478+
assert remote_server.recv(1024) == b"SELECT 1;"
479+
480+
# Step 2: Server sends response
481+
remote_server.sendall(b"RESULT 1")
482+
assert local_client.recv(1024) == b"RESULT 1"
483+
484+
# Step 3: Client closes connection while server is open
485+
local_client.close()
486+
487+
# Single-threaded proxy terminates immediately without deadlocking
488+
t_proxy.join(timeout=2.0)
489+
assert (
490+
not t_proxy.is_alive()
491+
), "Cloud SQL proxy must terminate cleanly without deadlocks"
492+
493+
remote_server.close()
494+
495+
496+
def test_proxy_pending_oserror_in_loop() -> None:
497+
"""Test that _proxy pending loop handles OSError on pending() and continues."""
498+
local_client, local_proxy = mockable_socketpair()
499+
remote_proxy, remote_server = mockable_socketpair()
500+
501+
# Add pending mock to real socket. First call returns 10, second raises OSError.
502+
remote_proxy.pending = MagicMock(side_effect=[10, OSError("pending failed")])
503+
504+
# Send data to be read by recv in pending loop
505+
remote_server.sendall(b"pending data")
506+
507+
# Start proxy in thread because it will block on select() after pending fails
508+
t = threading.Thread(target=_proxy, args=(local_proxy, remote_proxy), daemon=True)
509+
t.start()
510+
511+
# Verify local_client received the pending data
512+
assert local_client.recv(1024) == b"pending data"
513+
514+
# Now trigger proxy exit by closing local_client
515+
local_client.close()
516+
517+
t.join(timeout=2.0)
518+
assert not t.is_alive()
519+
520+
remote_server.close()
521+
522+
523+
def test_proxy_pending_non_int() -> None:
524+
"""Test that _proxy pending check handles non-int return values from pending()."""
525+
local_client, local_proxy = mockable_socketpair()
526+
remote_proxy, remote_server = mockable_socketpair()
527+
528+
# First call returns non-int. Should return False (goes to select)
529+
remote_proxy.pending = MagicMock(return_value="not an int")
530+
531+
# We need to run it in a thread because it will block on select
532+
t = threading.Thread(target=_proxy, args=(local_proxy, remote_proxy), daemon=True)
533+
t.start()
534+
535+
# Trigger exit.
536+
local_client.close()
537+
t.join(timeout=2.0)
538+
assert not t.is_alive()
539+
remote_server.close()
540+
541+
542+
def test_proxy_pending_non_int_in_loop() -> None:
543+
"""Test that _proxy pending loop handles non-int return values from pending() in loop."""
544+
local_client, local_proxy = mockable_socketpair()
545+
remote_proxy, remote_server = mockable_socketpair()
546+
547+
# Second call returns non-int.
548+
remote_proxy.pending = MagicMock(side_effect=[10, "not an int"])
549+
remote_server.sendall(b"data")
550+
551+
t = threading.Thread(target=_proxy, args=(local_proxy, remote_proxy), daemon=True)
552+
t.start()
553+
554+
assert local_client.recv(1024) == b"data"
555+
556+
local_client.close()
557+
t.join(timeout=2.0)
558+
assert not t.is_alive()
559+
remote_server.close()
560+
561+
562+
def test_proxy_pending_recv_eof() -> None:
563+
"""Test that _proxy pending loop handles remote recv EOF."""
564+
local_client, local_server = mockable_socketpair()
565+
remote_client, remote_server = mockable_socketpair()
566+
567+
remote_client.pending = MagicMock(return_value=10)
568+
# Mock recv to return EOF
569+
remote_client.recv = MagicMock(return_value=b"")
570+
571+
# Run proxy. It should call forward_pending, which calls recv, gets EOF,
572+
# and returns True. This breaks the loop and proxy exits.
573+
_proxy(local_server, remote_client)
574+
575+
# Sockets should be closed
576+
assert local_server.fileno() == -1
577+
assert remote_client.fileno() == -1
578+
579+
local_client.close()
580+
remote_server.close()
581+
582+
583+
def test_proxy_select_oserror() -> None:
584+
"""Test that _proxy loop handles OSError in select()."""
585+
local_client, local_server = mockable_socketpair()
586+
remote_client, remote_server = mockable_socketpair()
587+
588+
with patch(
589+
"google.cloud.sql.connector.psycopg.selectors.DefaultSelector"
590+
) as mock_sel_cls:
591+
mock_sel = MagicMock()
592+
mock_sel.select.side_effect = OSError("select failed")
593+
mock_sel_cls.return_value = mock_sel
594+
595+
# Run proxy. It should catch OSError on select() and exit loop to finally block.
596+
_proxy(local_server, remote_client)
597+
598+
assert local_server.fileno() == -1
599+
assert remote_client.fileno() == -1
600+
601+
local_client.close()
602+
remote_server.close()
603+
604+
605+
def test_proxy_finally_cleanup_errors() -> None:
606+
"""Test that _proxy finally block ignores OSErrors during socket shutdown/close."""
607+
local_client, local_server = mockable_socketpair()
608+
remote_client, remote_server = mockable_socketpair()
609+
610+
real_local_shutdown = local_server.shutdown
611+
def mock_local_shutdown(how):
612+
try:
613+
real_local_shutdown(how)
614+
except OSError:
615+
pass
616+
raise OSError("shutdown failed")
617+
local_server.shutdown = MagicMock(side_effect=mock_local_shutdown)
618+
619+
real_local_close = local_server.close
620+
def mock_local_close():
621+
real_local_close()
622+
raise OSError("close failed")
623+
local_server.close = MagicMock(side_effect=mock_local_close)
624+
625+
real_remote_shutdown = remote_client.shutdown
626+
def mock_remote_shutdown(how):
627+
try:
628+
real_remote_shutdown(how)
629+
except OSError:
630+
pass
631+
raise OSError("shutdown failed")
632+
remote_client.shutdown = MagicMock(side_effect=mock_remote_shutdown)
633+
634+
real_remote_close = remote_client.close
635+
def mock_remote_close():
636+
real_remote_close()
637+
raise OSError("close failed")
638+
remote_client.close = MagicMock(side_effect=mock_remote_close)
639+
640+
# Trigger exit by closing client
641+
local_client.close()
642+
643+
# Run proxy. It should handle the exceptions in finally block.
644+
_proxy(local_server, remote_client)
645+
646+
remote_server.close()
647+
648+
649+
@patch("google.cloud.sql.connector.psycopg._proxy")
650+
def test_accept_and_proxy_cleanup_errors(mock_proxy_fn: MagicMock) -> None:
651+
"""Test that accept_and_proxy thread handles errors during cleanup on exception."""
652+
# Make _proxy raise an exception to trigger the except block in _accept_and_proxy
653+
mock_proxy_fn.side_effect = Exception("proxy crash")
654+
655+
mock_unix_conn = MagicMock(spec=socket.socket)
656+
mock_unix_conn.shutdown.side_effect = OSError("unix conn shutdown failed")
657+
mock_unix_conn.close.side_effect = OSError("unix conn close failed")
658+
659+
mock_remote_sock = MagicMock(spec=ssl.SSLSocket)
660+
mock_remote_sock.shutdown.side_effect = OSError("remote shutdown failed")
661+
662+
# Mock socket.socket to return a mock local_sock that returns our mock_unix_conn
663+
real_socket = socket.socket
664+
mock_local_sock = MagicMock()
665+
mock_local_sock.accept.return_value = (mock_unix_conn, ("path",))
666+
667+
def socket_side_effect(family, type, proto=0, fileno=None):
668+
if family == socket.AF_UNIX:
669+
return mock_local_sock
670+
return real_socket(family, type, proto, fileno)
671+
672+
event = threading.Event()
673+
def remote_close_fn():
674+
event.set()
675+
raise OSError("remote close failed")
676+
mock_remote_sock.close.side_effect = remote_close_fn
677+
678+
with patch("socket.socket", side_effect=socket_side_effect), patch(
679+
"psycopg.connect"
680+
) as mock_psycopg_connect:
681+
mock_psycopg_connect.return_value = MagicMock()
682+
683+
connect(
684+
"127.0.0.1",
685+
mock_remote_sock,
686+
user="test_user",
687+
db="test_db",
688+
password="test_password",
689+
)
690+
691+
# Wait for the accept_and_proxy thread to finish cleanup
692+
assert event.wait(timeout=2.0)
693+
694+
# Verify mock_unix_conn shutdown and close were called
695+
mock_unix_conn.shutdown.assert_called_once_with(socket.SHUT_RDWR)
696+
mock_unix_conn.close.assert_called_once()
697+
mock_remote_sock.shutdown.assert_called_once_with(socket.SHUT_RDWR)
698+
mock_remote_sock.close.assert_called_once()
699+
700+
701+
@patch("psycopg.connect")
702+
def test_connect_wrapper_failure_cleanup_errors(mock_psycopg_connect: MagicMock) -> None:
703+
"""Test that connect wrapper ignores OSErrors during cleanup on connection failure."""
704+
mock_remote_sock = MagicMock(spec=ssl.SSLSocket)
705+
mock_remote_sock.close.side_effect = OSError("remote close failed")
706+
707+
mock_psycopg_connect.side_effect = Exception("connection failed simulated")
708+
709+
mock_local_sock = MagicMock()
710+
mock_local_sock.close.side_effect = OSError("local close failed")
711+
712+
real_socket = socket.socket
713+
def socket_side_effect(family, type, proto=0, fileno=None):
714+
if family == socket.AF_UNIX:
715+
return mock_local_sock
716+
return real_socket(family, type, proto, fileno)
717+
718+
with (
719+
patch("socket.socket", side_effect=socket_side_effect),
720+
pytest.raises(Exception, match="connection failed simulated"),
721+
):
722+
connect(
723+
"127.0.0.1",
724+
mock_remote_sock,
725+
user="test_user",
726+
db="test_db",
727+
password="test_password",
728+
)
729+
730+
mock_local_sock.close.assert_called_once()
731+
assert mock_remote_sock.close.call_count >= 1
732+
733+

0 commit comments

Comments
 (0)