Fix websocket.disconnect not delivered after server-initiated close

Signed-off-by: JaeHyuck Sa <jaehyuck.sa.dev@gmail.com>

Author: JaeHyuckSa

src/asgi/io.rs

@@ -434,13 +434,12 @@ impl ASGIWebsocketProtocol {
     #[inline(always)]
     fn close<'p>(&self, py: Python<'p>, frame: Option<wsframe::CloseFrame>) -> PyResult<Bound<'p, PyAny>> {
         let closed = self.closed.clone();
-        let ws_rx = self.ws_rx.clone();
         let ws_tx = self.ws_tx.clone();
 
         future_into_py_futlike(self.rt.clone(), py, async move {
             if let Some(tx) = ws_tx.lock().await.take() {
                 closed.store(true, atomic::Ordering::Release);
-                WebsocketDetachedTransport::new(true, ws_rx.lock().await.take(), Some(tx), frame)
+                WebsocketDetachedTransport::new(true, None, Some(tx), frame)
                     .close()
                     .await;
             }

tests/apps/asgi.py

@@ -1,6 +1,7 @@
 import asyncio
 import json
 import pathlib
+import tempfile
 
 import sniffio
 
@@ -109,6 +110,15 @@ async def ws_echo(scope, receive, send):
     await send({'type': 'websocket.close'})
 
 
+async def ws_server_close(scope, receive, send):
+    await receive()
+    await send({'type': 'websocket.accept'})
+    await receive()
+    await send({'type': 'websocket.close', 'code': 1000})
+    await asyncio.wait_for(receive(), timeout=5)
+    pathlib.Path(tempfile.gettempdir(), 'granian_ws_test_result').touch()
+
+
 async def ws_push(scope, receive, send):
     await send({'type': 'websocket.accept'})
 
@@ -188,6 +198,7 @@ def app(scope, receive, send):
         '/ws_reject': ws_reject,
         '/ws_info': ws_info,
         '/ws_echo': ws_echo,
+        '/ws_server_close': ws_server_close,
         '/ws_push': ws_push,
         '/err_app': err_app,
         '/err_proto/type': err_proto_msg,

tests/test_ws.py

@@ -1,5 +1,8 @@
+import asyncio
 import json
 import os
+import pathlib
+import tempfile
 
 import pytest
 import websockets
@@ -32,6 +35,28 @@ async def test_reject(server, runtime_mode):
     assert exc.value.response.status_code == 403
 
 
+@pytest.mark.asyncio
+@pytest.mark.parametrize('runtime_mode', ['mt', 'st'])
+async def test_server_initiated_close(asgi_server, runtime_mode):
+    result_path = pathlib.Path(tempfile.gettempdir(), 'granian_ws_test_result')
+    result_path.unlink(missing_ok=True)
+
+    async with asgi_server(runtime_mode) as port:
+        ws = await websockets.connect(f'ws://localhost:{port}/ws_server_close')
+        await ws.send('hello')
+        for _ in range(50):
+            if result_path.exists():
+                break
+            await asyncio.sleep(0.1)
+        try:
+            await ws.close()
+        except Exception:
+            pass
+
+    assert result_path.exists(), 'Server did not write result file'
+    result_path.unlink(missing_ok=True)
+
+
 @pytest.mark.asyncio
 @pytest.mark.skipif(bool(os.getenv('PGO_RUN')), reason='PGO build')
 @pytest.mark.parametrize('runtime_mode', ['mt', 'st'])