WsPool refactoring, TCP fallback backoff, removed direct connections outside of a pool

This commit is contained in:
Flowseal committed 2026-10-04 21:53:26 +03:00
1 parent 48954ff9ab
commit 50c980e9ea
13 files changed
+563 -294

No files matched your search

+24 -3
View File
@@ -19,8 +19,6 @@ class H2RoutingTest(unittest.IsolatedAsyncioTestCase):
with patch.object(proxy_config, 'dc_redirects', {2: '149.154.167.51'}), \
patch.object(proxy_config, 'force_test_dc', False), \
patch.object(tg_ws_proxy, 'cf_h2_pool', pool), \
patch.object(tg_ws_proxy, 'ws_blacklist', set()), \
patch.object(tg_ws_proxy, 'ip_fail_until', {}), \
patch.object(tg_ws_proxy, 'set_sock_opts'), \
patch.object(tg_ws_proxy, '_read_client_init', AsyncMock(
return_value=(b'', reader, writer, 'test'))), \
@@ -28,7 +26,6 @@ class H2RoutingTest(unittest.IsolatedAsyncioTestCase):
2, True, PROTO_TAG_INTERMEDIATE, b'')), \
patch.object(tg_ws_proxy, '_build_crypto_ctx'), \
patch.object(tg_ws_proxy.ws_pool, 'get', AsyncMock(return_value=ws)), \
patch.object(tg_ws_proxy.ws_pool, 'report_success'), \
patch.object(tg_ws_proxy, 'bridge_ws_reencrypt', AsyncMock()) as direct, \
patch.object(tg_ws_proxy, 'do_fallback', AsyncMock()) as fallback:
await tg_ws_proxy._handle_client(reader, writer, b'')
@@ -36,6 +33,30 @@ class H2RoutingTest(unittest.IsolatedAsyncioTestCase):
fallback.assert_not_awaited()
pool.open.assert_not_awaited()
async def test_pool_miss_goes_straight_to_fallback_including_test_dcs(self):
for dc, forced, expected_test in [(2, False, False),
(10002, False, True),
(2, True, True)]:
with self.subTest(dc=dc, forced=forced):
reader = Mock()
writer = Mock(wait_closed=AsyncMock())
writer.get_extra_info.return_value = None
with patch.object(proxy_config, 'force_test_dc', forced), \
patch.object(tg_ws_proxy, 'set_sock_opts'), \
patch.object(tg_ws_proxy, '_read_client_init', AsyncMock(
return_value=(b'', reader, writer, 'test'))), \
patch.object(tg_ws_proxy, '_try_handshake', return_value=(
dc, False, PROTO_TAG_INTERMEDIATE, b'')), \
patch.object(tg_ws_proxy, '_build_crypto_ctx'), \
patch.object(tg_ws_proxy.ws_pool, 'get', AsyncMock(return_value=None)) as get, \
patch('proxy.raw_websocket.RawWebSocket.connect', AsyncMock()) as connect, \
patch.object(tg_ws_proxy, 'do_fallback', AsyncMock(return_value=True)) as fallback:
await tg_ws_proxy._handle_client(reader, writer, b'')
get.assert_awaited_once_with(2, False, is_test_dc=expected_test)
fallback.assert_awaited_once()
self.assertEqual(fallback.call_args.args[4:7], (2, expected_test, False))
connect.assert_not_awaited()
async def test_cf_fallback_order_and_route_guards(self):
cases = [
# media, test DC, CF enabled, Worker succeeds, H2 opens, expected