From 53b4d44f55cca66ac33dab092ef2a30b1164b684 Mon Sep 17 00:00:00 2001
From: coletdjnz <coletdjnz@protonmail.com>
Date: Sat, 18 May 2024 19:12:21 +1200
Subject: [PATCH] [test] Fix connect timeout test (#9906)

Fixes https://github.com/yt-dlp/yt-dlp/issues/9659

Authored by: coletdjnz
---
 test/test_networking.py     | 23 ++++++++++-------------
 test/test_websockets.py     | 20 ++++++++++++++++++--
 yt_dlp/networking/common.py |  4 +++-
 3 files changed, 31 insertions(+), 16 deletions(-)

diff --git a/test/test_networking.py b/test/test_networking.py
index 994467014d..d127cbb94c 100644
--- a/test/test_networking.py
+++ b/test/test_networking.py
@@ -6,7 +6,7 @@ import sys
 
 import pytest
 
-from yt_dlp.networking.common import Features
+from yt_dlp.networking.common import Features, DEFAULT_TIMEOUT
 
 sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
 
@@ -523,20 +523,17 @@ class TestHTTPRequestHandler(TestRequestHandlerBase):
     def test_connect_timeout(self, handler):
         # nothing should be listening on this port
         connect_timeout_url = 'http://10.255.255.255'
-        with handler(timeout=0.01) as rh:
+        with handler(timeout=0.01) as rh, pytest.raises(TransportError):
             now = time.time()
-            with pytest.raises(TransportError):
-                validate_and_send(
-                    rh, Request(connect_timeout_url))
-            assert 0.01 <= time.time() - now < 20
+            validate_and_send(rh, Request(connect_timeout_url))
+        assert time.time() - now < DEFAULT_TIMEOUT
 
-        with handler() as rh:
-            with pytest.raises(TransportError):
-                # Per request timeout, should override handler timeout
-                now = time.time()
-                validate_and_send(
-                    rh, Request(connect_timeout_url, extensions={'timeout': 0.01}))
-                assert 0.01 <= time.time() - now < 20
+        # Per request timeout, should override handler timeout
+        request = Request(connect_timeout_url, extensions={'timeout': 0.01})
+        with handler() as rh, pytest.raises(TransportError):
+            now = time.time()
+            validate_and_send(rh, request)
+        assert time.time() - now < DEFAULT_TIMEOUT
 
     def test_source_address(self, handler):
         source_address = f'127.0.0.{random.randint(5, 255)}'
diff --git a/test/test_websockets.py b/test/test_websockets.py
index bc9f2187a1..aa0dfa2d55 100644
--- a/test/test_websockets.py
+++ b/test/test_websockets.py
@@ -3,11 +3,12 @@
 # Allow direct execution
 import os
 import sys
+import time
 
 import pytest
 
 from test.helper import verify_address_availability
-from yt_dlp.networking.common import Features
+from yt_dlp.networking.common import Features, DEFAULT_TIMEOUT
 
 sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
 
@@ -202,11 +203,26 @@ class TestWebsSocketRequestHandlerConformance:
         ({'timeout': sys.float_info.min}, {}),
         ({}, {'timeout': sys.float_info.min}),
     ])
-    def test_timeout(self, handler, params, extensions):
+    def test_read_timeout(self, handler, params, extensions):
         with handler(**params) as rh:
             with pytest.raises(TransportError):
                 ws_validate_and_send(rh, Request(self.ws_base_url, extensions=extensions))
 
+    def test_connect_timeout(self, handler):
+        # nothing should be listening on this port
+        connect_timeout_url = 'ws://10.255.255.255'
+        with handler(timeout=0.01) as rh, pytest.raises(TransportError):
+            now = time.time()
+            ws_validate_and_send(rh, Request(connect_timeout_url))
+        assert time.time() - now < DEFAULT_TIMEOUT
+
+        # Per request timeout, should override handler timeout
+        request = Request(connect_timeout_url, extensions={'timeout': 0.01})
+        with handler() as rh, pytest.raises(TransportError):
+            now = time.time()
+            ws_validate_and_send(rh, request)
+        assert time.time() - now < DEFAULT_TIMEOUT
+
     def test_cookies(self, handler):
         cookiejar = YoutubeDLCookieJar()
         cookiejar.set_cookie(http.cookiejar.Cookie(
diff --git a/yt_dlp/networking/common.py b/yt_dlp/networking/common.py
index a2217034c9..d473e16c5a 100644
--- a/yt_dlp/networking/common.py
+++ b/yt_dlp/networking/common.py
@@ -31,6 +31,8 @@ from ..utils import (
 )
 from ..utils.networking import HTTPHeaderDict, normalize_url
 
+DEFAULT_TIMEOUT = 20
+
 
 def register_preference(*handlers: type[RequestHandler]):
     assert all(issubclass(handler, RequestHandler) for handler in handlers)
@@ -235,7 +237,7 @@ class RequestHandler(abc.ABC):
         self._logger = logger
         self.headers = headers or {}
         self.cookiejar = cookiejar if cookiejar is not None else YoutubeDLCookieJar()
-        self.timeout = float(timeout or 20)
+        self.timeout = float(timeout or DEFAULT_TIMEOUT)
         self.proxies = proxies or {}
         self.source_address = source_address
         self.verbose = verbose