Merging upstream version 2024.05.26.

Signed-off-by: Daniel Baumann <daniel.baumann@progress-linux.org>

1 files changed, 58 insertions, 15 deletions

diff --git a/test/test_websockets.py b/test/test_websockets.py
index b294b09..aa0dfa2 100644
--- a/test/test_websockets.py
+++ b/test/test_websockets.py
@@ -3,10 +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, DEFAULT_TIMEOUT
 
 sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
 
@@ -18,7 +20,7 @@ import random
 import ssl
 import threading
 
-from yt_dlp import socks
+from yt_dlp import socks, traverse_obj
 from yt_dlp.cookies import YoutubeDLCookieJar
 from yt_dlp.dependencies import websockets
 from yt_dlp.networking import Request
@@ -114,6 +116,7 @@ def ws_validate_and_send(rh, req):
 
 
 @pytest.mark.skipif(not websockets, reason='websockets must be installed to test websocket request handlers')
+@pytest.mark.parametrize('handler', ['Websockets'], indirect=True)
 class TestWebsSocketRequestHandlerConformance:
     @classmethod
     def setup_class(cls):
@@ -129,7 +132,6 @@ class TestWebsSocketRequestHandlerConformance:
         cls.mtls_wss_thread, cls.mtls_wss_port = create_mtls_wss_websocket_server()
         cls.mtls_wss_base_url = f'wss://127.0.0.1:{cls.mtls_wss_port}'
 
-    @pytest.mark.parametrize('handler', ['Websockets'], indirect=True)
     def test_basic_websockets(self, handler):
         with handler() as rh:
             ws = ws_validate_and_send(rh, Request(self.ws_base_url))
@@ -141,7 +143,6 @@ class TestWebsSocketRequestHandlerConformance:
 
     # https://www.rfc-editor.org/rfc/rfc6455.html#section-5.6
     @pytest.mark.parametrize('msg,opcode', [('str', 1), (b'bytes', 2)])
-    @pytest.mark.parametrize('handler', ['Websockets'], indirect=True)
     def test_send_types(self, handler, msg, opcode):
         with handler() as rh:
             ws = ws_validate_and_send(rh, Request(self.ws_base_url))
@@ -149,7 +150,6 @@ class TestWebsSocketRequestHandlerConformance:
             assert int(ws.recv()) == opcode
             ws.close()
 
-    @pytest.mark.parametrize('handler', ['Websockets'], indirect=True)
     def test_verify_cert(self, handler):
         with handler() as rh:
             with pytest.raises(CertificateVerifyError):
@@ -160,14 +160,12 @@ class TestWebsSocketRequestHandlerConformance:
             assert ws.status == 101
             ws.close()
 
-    @pytest.mark.parametrize('handler', ['Websockets'], indirect=True)
     def test_ssl_error(self, handler):
         with handler(verify=False) as rh:
             with pytest.raises(SSLError, match=r'ssl(?:v3|/tls) alert handshake failure') as exc_info:
                 ws_validate_and_send(rh, Request(self.bad_wss_host))
             assert not issubclass(exc_info.type, CertificateVerifyError)
 
-    @pytest.mark.parametrize('handler', ['Websockets'], indirect=True)
     @pytest.mark.parametrize('path,expected', [
         # Unicode characters should be encoded with uppercase percent-encoding
         ('/中文', '/%E4%B8%AD%E6%96%87'),
@@ -182,7 +180,6 @@ class TestWebsSocketRequestHandlerConformance:
             assert ws.status == 101
             ws.close()
 
-    @pytest.mark.parametrize('handler', ['Websockets'], indirect=True)
     def test_remove_dot_segments(self, handler):
         with handler() as rh:
             # This isn't a comprehensive test,
@@ -195,7 +192,6 @@ class TestWebsSocketRequestHandlerConformance:
 
     # We are restricted to known HTTP status codes in http.HTTPStatus
     # Redirects are not supported for websockets
-    @pytest.mark.parametrize('handler', ['Websockets'], indirect=True)
     @pytest.mark.parametrize('status', (200, 204, 301, 302, 303, 400, 500, 511))
     def test_raise_http_error(self, handler, status):
         with handler() as rh:
@@ -203,17 +199,30 @@ class TestWebsSocketRequestHandlerConformance:
                 ws_validate_and_send(rh, Request(f'{self.ws_base_url}/gen_{status}'))
             assert exc_info.value.status == status
 
-    @pytest.mark.parametrize('handler', ['Websockets'], indirect=True)
     @pytest.mark.parametrize('params,extensions', [
         ({'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))
 
-    @pytest.mark.parametrize('handler', ['Websockets'], indirect=True)
+    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(
@@ -239,7 +248,6 @@ class TestWebsSocketRequestHandlerConformance:
             assert json.loads(ws.recv())['cookie'] == 'test=ytdlp'
             ws.close()
 
-    @pytest.mark.parametrize('handler', ['Websockets'], indirect=True)
     def test_source_address(self, handler):
         source_address = f'127.0.0.{random.randint(5, 255)}'
         verify_address_availability(source_address)
@@ -249,7 +257,6 @@ class TestWebsSocketRequestHandlerConformance:
             assert source_address == ws.recv()
             ws.close()
 
-    @pytest.mark.parametrize('handler', ['Websockets'], indirect=True)
     def test_response_url(self, handler):
         with handler() as rh:
             url = f'{self.ws_base_url}/something'
@@ -257,7 +264,6 @@ class TestWebsSocketRequestHandlerConformance:
             assert ws.url == url
             ws.close()
 
-    @pytest.mark.parametrize('handler', ['Websockets'], indirect=True)
     def test_request_headers(self, handler):
         with handler(headers=HTTPHeaderDict({'test1': 'test', 'test2': 'test2'})) as rh:
             # Global Headers
@@ -293,7 +299,6 @@ class TestWebsSocketRequestHandlerConformance:
             'client_certificate_password': 'foobar',
         }
     ))
-    @pytest.mark.parametrize('handler', ['Websockets'], indirect=True)
     def test_mtls(self, handler, client_cert):
         with handler(
             # Disable client-side validation of unacceptable self-signed testcert.pem
@@ -303,6 +308,44 @@ class TestWebsSocketRequestHandlerConformance:
         ) as rh:
             ws_validate_and_send(rh, Request(self.mtls_wss_base_url)).close()
 
+    def test_request_disable_proxy(self, handler):
+        for proxy_proto in handler._SUPPORTED_PROXY_SCHEMES or ['ws']:
+            # Given handler is configured with a proxy
+            with handler(proxies={'ws': f'{proxy_proto}://10.255.255.255'}, timeout=5) as rh:
+                # When a proxy is explicitly set to None for the request
+                ws = ws_validate_and_send(rh, Request(self.ws_base_url, proxies={'http': None}))
+                # Then no proxy should be used
+                assert ws.status == 101
+                ws.close()
+
+    @pytest.mark.skip_handlers_if(
+        lambda _, handler: Features.NO_PROXY not in handler._SUPPORTED_FEATURES, 'handler does not support NO_PROXY')
+    def test_noproxy(self, handler):
+        for proxy_proto in handler._SUPPORTED_PROXY_SCHEMES or ['ws']:
+            # Given the handler is configured with a proxy
+            with handler(proxies={'ws': f'{proxy_proto}://10.255.255.255'}, timeout=5) as rh:
+                for no_proxy in (f'127.0.0.1:{self.ws_port}', '127.0.0.1', 'localhost'):
+                    # When request no proxy includes the request url host
+                    ws = ws_validate_and_send(rh, Request(self.ws_base_url, proxies={'no': no_proxy}))
+                    # Then the proxy should not be used
+                    assert ws.status == 101
+                    ws.close()
+
+    @pytest.mark.skip_handlers_if(
+        lambda _, handler: Features.ALL_PROXY not in handler._SUPPORTED_FEATURES, 'handler does not support ALL_PROXY')
+    def test_allproxy(self, handler):
+        supported_proto = traverse_obj(handler._SUPPORTED_PROXY_SCHEMES, 0, default='ws')
+        # This is a bit of a hacky test, but it should be enough to check whether the handler is using the proxy.
+        # 0.1s might not be enough of a timeout if proxy is not used in all cases, but should still get failures.
+        with handler(proxies={'all': f'{supported_proto}://10.255.255.255'}, timeout=0.1) as rh:
+            with pytest.raises(TransportError):
+                ws_validate_and_send(rh, Request(self.ws_base_url)).close()
+
+        with handler(timeout=0.1) as rh:
+            with pytest.raises(TransportError):
+                ws_validate_and_send(
+                    rh, Request(self.ws_base_url, proxies={'all': f'{supported_proto}://10.255.255.255'})).close()
+
 
 def create_fake_ws_connection(raised):
     import websockets.sync.client