Repository: yt-dlp/yt-dlp
File: test/test_websockets.py (Python)

Code before:

MTLS_CERT_DIR = os.path.join(TEST_DIR, 'testdata', 'certificate')


def create_mtls_wss_websocket_server():
    certfn = os.path.join(TEST_DIR, 'testcert.pem')
    cacertfn = os.path.join(MTLS_CERT_DIR, 'ca.crt')

    sslctx = ssl.SSLContext(ssl.PROTOCOL_TLS_SERVER)
    sslctx.verify_mode = ssl.CERT_REQUIRED
    sslctx.load_verify_locations(cafile=cacertfn)
    sslctx.load_cert_chain(certfn, None)

    return create_websocket_server(ssl_context=sslctx)


def ws_validate_and_send(rh, req):
    rh.validate(req)
    for i in range(0, 3):
        try:
            return rh.send(req)
        except TransportError as e:
            if 'connection closed during handshake' in str(e):
                # websockets server sometimes hangs on new connections
                continue
            raise


@pytest.mark.skipif(not websockets, reason='websockets must be installed to test websocket request handlers')
class TestWebsSocketRequestHandlerConformance:
    @classmethod
    def setup_class(cls):
        cls.ws_thread, cls.ws_port = create_ws_websocket_server()
        cls.ws_base_url = f'ws://127.0.0.1:{cls.ws_port}'

        cls.wss_thread, cls.wss_port = create_wss_websocket_server()
        cls.wss_base_url = f'wss://127.0.0.1:{cls.wss_port}'

        cls.bad_wss_thread, cls.bad_wss_port = create_websocket_server(ssl_context=ssl.SSLContext(ssl.PROTOCOL_TLS_SERVER))
        cls.bad_wss_host = f'wss://127.0.0.1:{cls.bad_wss_port}'

        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))
            assert 'upgrade' in ws.headers
            assert ws.status == 101
            ws.send('foo')

Review comment (suggestion):
would be good to raise when retry loop is exhausted
```suggestion
    max_tries = 3
    for i in range(max_tries):
        try:
            return rh.send(req)
        except TransportError as e:
            if i < (max_tries - 1) and 'connection closed during handshake' in str(e):
                # websockets server sometimes hangs on new connections
                continue
            raise
```

Code after:
    return create_websocket_server(ssl_context=sslctx)


MTLS_CERT_DIR = os.path.join(TEST_DIR, 'testdata', 'certificate')


def create_mtls_wss_websocket_server():
    certfn = os.path.join(TEST_DIR, 'testcert.pem')
    cacertfn = os.path.join(MTLS_CERT_DIR, 'ca.crt')

    sslctx = ssl.SSLContext(ssl.PROTOCOL_TLS_SERVER)
    sslctx.verify_mode = ssl.CERT_REQUIRED
    sslctx.load_verify_locations(cafile=cacertfn)
    sslctx.load_cert_chain(certfn, None)

    return create_websocket_server(ssl_context=sslctx)


def ws_validate_and_send(rh, req):
    rh.validate(req)
    max_tries = 3
    for i in range(max_tries):
        try:
            return rh.send(req)
        except TransportError as e:
            if i < (max_tries - 1) and 'connection closed during handshake' in str(e):
                # websockets server sometimes hangs on new connections
                continue
            raise


@pytest.mark.skipif(not websockets, reason='websockets must be installed to test websocket request handlers')
class TestWebsSocketRequestHandlerConformance:
    @classmethod
    def setup_class(cls):
        cls.ws_thread, cls.ws_port = create_ws_websocket_server()
        cls.ws_base_url = f'ws://127.0.0.1:{cls.ws_port}'

        cls.wss_thread, cls.wss_port = create_wss_websocket_server()
        cls.wss_base_url = f'wss://127.0.0.1:{cls.wss_port}'

        cls.bad_wss_thread, cls.bad_wss_port = create_websocket_server(ssl_context=ssl.SSLContext(ssl.PROTOCOL_TLS_SERVER))
        cls.bad_wss_host = f'wss://127.0.0.1:{cls.bad_wss_port}'

        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))