Allow overriding the scheme used in websocket_connect to support wss (#…

…48)

async_asgi_testclient/testing.py

@@ -67,9 +67,9 @@ def __init__(
         scope: Optional[dict] = None,
     ):
         self.application = guarantee_single_callable(application)
-        self.cookie_jar: Optional[
-            SimpleCookie
-        ] = SimpleCookie() if use_cookies else None
+        self.cookie_jar: Optional[SimpleCookie] = (
+            SimpleCookie() if use_cookies else None
+        )
         self.timeout = timeout
         self.headers = headers or {}
         self._scope = scope or {}
@@ -106,8 +106,8 @@ async def send_lifespan(self, action):
         elif message["type"] == f"lifespan.{action}.failed":
             raise Exception(message)
 
-    def websocket_connect(self, path, headers=None, cookies=None):
-        return WebSocketSession(self, path, headers, cookies)
+    def websocket_connect(self, *args: Any, **kwargs: Any) -> WebSocketSession:
+        return WebSocketSession(self, *args, **kwargs)
 
     async def open(
         self,
@@ -304,41 +304,33 @@ async def wait_response(self, receive_or_fail, type_):
         return message
 
     async def delete(self, *args: Any, **kwargs: Any) -> Response:
-        """Make a DELETE request.
-        """
+        """Make a DELETE request."""
         return await self.open(*args, method="DELETE", **kwargs)
 
     async def get(self, *args: Any, **kwargs: Any) -> Response:
-        """Make a GET request.
-        """
+        """Make a GET request."""
         return await self.open(*args, method="GET", **kwargs)
 
     async def head(self, *args: Any, **kwargs: Any) -> Response:
-        """Make a HEAD request.
-        """
+        """Make a HEAD request."""
         return await self.open(*args, method="HEAD", **kwargs)
 
     async def options(self, *args: Any, **kwargs: Any) -> Response:
-        """Make a OPTIONS request.
-        """
+        """Make a OPTIONS request."""
         return await self.open(*args, method="OPTIONS", **kwargs)
 
     async def patch(self, *args: Any, **kwargs: Any) -> Response:
-        """Make a PATCH request.
-        """
+        """Make a PATCH request."""
         return await self.open(*args, method="PATCH", **kwargs)
 
     async def post(self, *args: Any, **kwargs: Any) -> Response:
-        """Make a POST request.
-        """
+        """Make a POST request."""
         return await self.open(*args, method="POST", **kwargs)
 
     async def put(self, *args: Any, **kwargs: Any) -> Response:
-        """Make a PUT request.
-        """
+        """Make a PUT request."""
         return await self.open(*args, method="PUT", **kwargs)
 
     async def trace(self, *args: Any, **kwargs: Any) -> Response:
-        """Make a TRACE request.
-        """
+        """Make a TRACE request."""
         return await self.open(*args, method="TRACE", **kwargs)

async_asgi_testclient/tests/test_testing.py

@@ -103,6 +103,8 @@ class Echo(WebSocketEndpoint):
         async def on_receive(self, websocket, data):
             if data == "cookies":
                 await websocket.send_text(dumps(websocket.cookies))
+            elif data == "url":
+                await websocket.send_text(str(websocket.url))
             else:
                 await websocket.send_text(f"Message text was: {data}")
 
@@ -430,6 +432,24 @@ async def test_ws_connect_inherits_test_client_cookies(starlette_app):
             assert msg == '{"session": "abc"}'
 
 
+@pytest.mark.asyncio
+async def test_ws_connect_default_scheme(starlette_app):
+    async with TestClient(starlette_app, timeout=0.1) as client:
+        async with client.websocket_connect("/ws") as ws:
+            await ws.send_text("url")
+            msg = await ws.receive_text()
+            assert msg.startswith("ws://")
+
+
+@pytest.mark.asyncio
+async def test_ws_connect_custom_scheme(starlette_app):
+    async with TestClient(starlette_app, timeout=0.1) as client:
+        async with client.websocket_connect("/ws", scheme="wss") as ws:
+            await ws.send_text("url")
+            msg = await ws.receive_text()
+            assert msg.startswith("wss://")
+
+
 @pytest.mark.asyncio
 async def test_request_stream(starlette_app):
     from starlette.responses import StreamingResponse

async_asgi_testclient/websocket.py

@@ -18,11 +18,13 @@ def __init__(
         path,
         headers: Optional[Dict] = None,
         cookies: Optional[Dict] = None,
+        scheme: str = "ws",
     ):
         self.testclient = testclient
         self.path = path
         self.headers = headers or {}
         self.cookies = cookies
+        self.scheme = scheme
         self.input_queue: asyncio.Queue[dict] = asyncio.Queue()
         self.output_queue: asyncio.Queue[dict] = asyncio.Queue()
         self._app_task = None  # Necessary to keep a hard reference to running task
@@ -118,7 +120,7 @@ async def connect(self):
             "path": path,
             "query_string": query_string_bytes,
             "root_path": "",
-            "scheme": "ws",
+            "scheme": self.scheme,
             "subprotocols": [],
         }