From 35caeb53d10d106affa54ab174eb4cb9ffd4967d Mon Sep 17 00:00:00 2001
From: Dmitry Kropachev <dmitry.kropachev@gmail.com>
Date: Fri, 31 Jul 2026 22:32:55 -0400
Subject: [PATCH] test: track TcpProxy forwarders through socket cleanup

---
 .../standard/test_client_routes.py            |  2 +-
 tests/unit/test_tcp_proxy.py                  | 39 ++++++++++++++++++-
 2 files changed, 39 insertions(+), 2 deletions(-)

diff --git a/tests/integration/standard/test_client_routes.py b/tests/integration/standard/test_client_routes.py
index 8e45cf7d93..5647d65829 100644
--- a/tests/integration/standard/test_client_routes.py
+++ b/tests/integration/standard/test_client_routes.py
@@ -224,9 +224,9 @@ def _forward_loop(self, client_sock, target_sock):
         except (OSError, ConnectionResetError, BrokenPipeError):
             pass
         finally:
+            self._close_pair(client_sock, target_sock)
             with self._lock:
                 self._connections.pop((client_sock, target_sock), None)
-            self._close_pair(client_sock, target_sock)
 
     @staticmethod
     def _close_pair(csock, tsock):
diff --git a/tests/unit/test_tcp_proxy.py b/tests/unit/test_tcp_proxy.py
index 71d47d6150..91bfdc6a4c 100644
--- a/tests/unit/test_tcp_proxy.py
+++ b/tests/unit/test_tcp_proxy.py
@@ -14,7 +14,7 @@
 
 """
 Regression tests for the ``TcpProxy`` test helper's connection
-shutdown/join synchronization path (GitHub issue #948).
+shutdown/join synchronization path (GitHub issues #948 and #962).
 
 ``TcpProxy`` is defined in
 ``tests/integration/standard/test_client_routes.py`` because it backs the
@@ -176,6 +176,43 @@ def test_timed_out_forwarder_thread_is_retained_until_it_exits(self):
         self.assertEqual(self.proxy.active_connections, 0)
         self.assertNotIn((csock, tsock), self.proxy._connections)
 
+    def test_forwarder_is_tracked_until_socket_cleanup_finishes(self):
+        """Keep a connection registered while its sockets are closing."""
+        client = _open_client(self.proxy.listen_host, self.proxy.listen_port)
+        self.addCleanup(client.close)
+        client.sendall(b"ping")
+        self.assertEqual(client.recv(16), b"ping")
+
+        self.assertEqual(self.proxy.active_connections, 1)
+        connection, thread = next(iter(self.proxy._connections.items()))
+        cleanup_started = threading.Event()
+        allow_cleanup = threading.Event()
+        real_close_pair = TcpProxy._close_pair
+
+        def blocking_close_pair(csock, tsock):
+            cleanup_started.set()
+            allow_cleanup.wait()
+            real_close_pair(csock, tsock)
+
+        try:
+            with patch.object(TcpProxy, "_close_pair",
+                              new=staticmethod(blocking_close_pair)):
+                client.shutdown(socket.SHUT_RDWR)
+                self.assertTrue(
+                    cleanup_started.wait(timeout=5),
+                    "forwarder did not begin socket cleanup")
+
+                self.assertTrue(thread.is_alive())
+                self.assertEqual(self.proxy.active_connections, 1)
+                self.assertIn(connection, self.proxy._connections)
+        finally:
+            allow_cleanup.set()
+
+        thread.join(timeout=5)
+        self.assertFalse(thread.is_alive())
+        self.assertEqual(self.proxy.active_connections, 0)
+        self.assertNotIn(connection, self.proxy._connections)
+
     def test_concurrent_stop_and_drop_leaves_no_live_forwarders(self):
         """
         Deterministic stress regression test: concurrently open/close real