Merge pull request #74 from Misutes/feature/websocket_correlation_id

add correlation_id into websocket request
snok · Apr 22, 2023 · 75dbf55 · 75dbf55
2 parents bf3b500 + 8ee1b4e
commit 75dbf55
Show file tree

Hide file tree

Showing 2 changed files with 4 additions and 1 deletion.
diff --git a/asgi_correlation_id/middleware.py b/asgi_correlation_id/middleware.py
@@ -46,7 +46,7 @@ async def __call__(self, scope: 'Scope', receive: 'Receive', send: 'Send') -> No
         """
         Load request ID from headers if present. Generate one otherwise.
         """
-        if scope['type'] != 'http':
+        if scope['type'] not in ('http', 'websocket'):
             await self.app(scope, receive, send)
             return
 

diff --git a/tests/test_middleware.py b/tests/test_middleware.py
@@ -120,6 +120,9 @@ async def test_websocket_request(caplog, app):
 
     @app.websocket_route('/ws')
     async def websocket(websocket: 'WebSocket'):
+        # Check we get the right headers back
+        assert websocket.headers.get('x-request-id') is not None
+
         await websocket.accept()
         await websocket.send_json({'msg': 'Hello WebSocket'})
         await websocket.close()