Repository: litestar-org/litestar (python)

Issue:
Docs: Incorrect example for websocket listener `connection_lifespan` 
### Description

I am trying the example , and it is not working. I had drill down to included MVCE and i believe  connection_lifespan is the culprit.

### URL to code causing the issue

https://github.com/litestar-org/litestar/blob/027ce12a64711f90c9d13b79d753be81779e6db0/docs/examples/websockets/stream_and_receive_listener.py

### MCVE

```python
import asyncio
import time
from typing import Any, AsyncGenerator

from litestar import Litestar, WebSocket, websocket_listener
from litestar.handlers import send_websocket_stream


async def listener_lifespan(socket: WebSocket) -> None:
    pass


@websocket_listener("/", connection_lifespan=listener_lifespan)
def handler(socket: WebSocket, data: Any) -> None:
    print(f"{socket.client}: {data}")


app = Litestar([handler])
```

### Steps to reproduce

1. litestar run --reload  -p 8881 --host=0.0.0.0 
2. websocat ws://localhost:8881


### Screenshots

_No response_

### Logs

```text
INFO:     ('127.0.0.1', 38680) - "WebSocket /" 403
INFO:     connection rejected (403 Forbidden)
INFO:     connection closed



websocat: WebSocketError: Connection refused (os error 111)
websocat: error running
(litestar-chat)
```

### Litestar Version

027ce12a64711f90c9d13b79d753be81779e6db0

### Platform

- [x] Linux
- [ ] Mac
- [ ] Windows
- [ ] Other (Please specify in the description above)

Pull request:
docs: add test to stream_and_receive_listener and mention on json default

<!--
By submitting this pull request, you agree to:
- follow [Litestar's Code of Conduct](https://github.com/litestar-org/.github/blob/main/CODE_OF_CONDUCT.md)
- follow [Litestar's contribution guidelines](https://github.com/litestar-org/.github/blob/main/CONTRIBUTING.md)
- follow the [PSFs's Code of Conduct](https://www.python.org/psf/conduct/)
-->
## Description

-

<!--
Please add in issue numbers this pull request will close, if applicable
Examples: Fixes #4321 or Closes #1234

Ensure you are using a supported keyword to properly link an issue:
https://docs.github.com/en/issues/tracking-your-work-with-issues/linking-a-pull-request-to-an-issue#linking-a-pull-request-to-an-issue-using-a-keyword
-->
## Closes

fixes #4100 


Patch:
diff --git a/docs/examples/websockets/stream_and_receive_listener.py b/docs/examples/websockets/stream_and_receive_listener.py
index d01158e12..5b37b59dc 100644
--- a/docs/examples/websockets/stream_and_receive_listener.py
+++ b/docs/examples/websockets/stream_and_receive_listener.py
@@ -1,26 +1,36 @@
-import asyncio
-import time
+from contextlib import asynccontextmanager
 from typing import Any, AsyncGenerator
 
+import anyio
+
 from litestar import Litestar, WebSocket, websocket_listener
+from litestar.exceptions import WebSocketDisconnect
 from litestar.handlers import send_websocket_stream
 
 
-async def listener_lifespan(socket: WebSocket) -> None:
-    async def handle_stream() -> AsyncGenerator[dict[str, float], None]:
-        while True:
-            yield {"time": time.time()}
-            await asyncio.sleep(0.5)
+@asynccontextmanager
+async def listener_lifespan(socket: WebSocket) -> AsyncGenerator[None, Any]:
+    is_closed = anyio.Event()
+
+    async def handle_stream() -> AsyncGenerator[str, None]:
+        while not is_closed.is_set():
+            await anyio.sleep(0.1)
+            yield "ping"
+
+    async with anyio.create_task_group() as tg:
+        tg.start_soon(send_websocket_stream, socket, handle_stream())
 
-    task = asyncio.create_task(send_websocket_stream(socket=socket, stream=handle_stream()))
-    yield
-    task.cancel()
-    await task
+        try:
+            yield
+        except WebSocketDisconnect:
+            pass
+        finally:
+            is_closed.set()
 
 
 @websocket_listener("/", connection_lifespan=listener_lifespan)
-def handler(socket: WebSocket, data: Any) -> None:
-    print(f"{socket.client}: {data}")
+def handler(data: str) -> str:
+    return data
 
 
 app = Litestar([handler])
diff --git a/docs/examples/websockets/stream_and_receive_raw.py b/docs/examples/websockets/stream_and_receive_raw.py
index 6e6966b0e..e4d895de3 100644
--- a/docs/examples/websockets/stream_and_receive_raw.py
+++ b/docs/examples/websockets/stream_and_receive_raw.py
@@ -1,27 +1,33 @@
-import asyncio
-import time
-from typing import AsyncGenerator
+from typing import Any, AsyncGenerator
+
+import anyio
 
 from litestar import Litestar, WebSocket, websocket
+from litestar.exceptions import WebSocketDisconnect
 from litestar.handlers import send_websocket_stream
 
 
 @websocket("/")
 async def handler(socket: WebSocket) -> None:
     await socket.accept()
+    should_stop = anyio.Event()
 
-    async def handle_stream() -> AsyncGenerator[dict[str, float], None]:
-        while True:
-            yield {"time": time.time()}
-            await asyncio.sleep(0.5)
+    async def handle_stream() -> AsyncGenerator[str, None]:
+        while not should_stop.is_set():
+            await anyio.sleep(0.5)
+            yield "ping"
 
-    async def handle_receive() -> None:
+    async def handle_receive() -> Any:
+        await socket.send_json({"handle_receive": "start"})
         async for event in socket.iter_json():
-            print(f"{socket.client}: {event}")
-
-    async with asyncio.TaskGroup() as tg:
-        tg.create_task(send_websocket_stream(socket=socket, stream=handle_stream()))
-        tg.create_task(handle_receive())
+            await socket.send_json(event)
+
+    try:
+        async with anyio.create_task_group() as tg:
+            tg.start_soon(send_websocket_stream, socket, handle_stream())
+            tg.start_soon(handle_receive)
+    except WebSocketDisconnect:
+        should_stop.set()
 
 
 app = Litestar([handler])
diff --git a/docs/usage/websockets.rst b/docs/usage/websockets.rst
index cce14f013..d52913e46 100644
--- a/docs/usage/websockets.rst
+++ b/docs/usage/websockets.rst
@@ -61,8 +61,12 @@ Receiving data
 
 Data can be received in the listener via the ``data`` parameter. The data passed to this
 will be converted / parsed according to the given type annotation and supports
-:class:`str`, :class:`bytes`, or arbitrary :class:`dict`\ s / or :class:`list`\ s in the
+:class:`str`, :class:`bytes`, or arbitrary :class:`dict` / or :class:`list` in the
 form of JSON.
+.. important::
+
+    The listeners will default to JSON unless `data` is annotated with `str` or `bytes`
+
 
 .. tab-set::
 
@@ -365,14 +369,34 @@ To stream and receive data concurrently, the stream can be set up manually using
 
     .. tab-item:: websocket_listener
 
-        .. literalinclude:: /examples/websockets/stream_and_receive_listener.py
-            :language: python
+        .. tab-set::
+
+            .. tab-item:: example
+
+                .. literalinclude:: /examples/websockets/stream_and_receive_listener.py
+                    :language: python
+
+            .. tab-item:: how to test
+
+                .. literalinclude:: ../../tests/examples/test_websockets.py
+                    :language: python
+                    :lines: 18-25
+
 
     .. tab-item:: websocket handler
 
-        .. literalinclude:: /examples/websockets/stream_and_receive_raw.py
-            :language: python
+        .. tab-set::
+
+            .. tab-item:: example
+
+                .. literalinclude:: /examples/websockets/stream_and_receive_raw.py
+                    :language: python
+
+            .. tab-item:: how to test
 
+                .. literalinclude:: ../../tests/examples/test_websockets.py
+                    :language: python
+                    :lines: 28-35
 
 Transport modes
 ---------------


Test patch:
diff --git a/tests/examples/test_websockets.py b/tests/examples/test_websockets.py
index c2898aa11..0599c26fc 100644
--- a/tests/examples/test_websockets.py
+++ b/tests/examples/test_websockets.py
@@ -1,5 +1,8 @@
 from docs.examples.websockets.custom_websocket import app as custom_websocket_class_app
+from docs.examples.websockets.stream_and_receive_listener import app as app_stream_and_receive_listener
+from docs.examples.websockets.stream_and_receive_raw import app as app_stream_and_receive_raw
 
+from litestar.testing import AsyncTestClient
 from litestar.testing.client.sync_client import TestClient
 
 
@@ -10,3 +13,23 @@ def test_custom_websocket_class():
         ws.send({"data": "I should not be in response"})
         data = ws.receive()
         assert data["text"] == "Fixed response"
+
+
+async def test_websocket_listener() -> None:
+    """Test the websocket listener."""
+    async with AsyncTestClient(app_stream_and_receive_listener) as client:
+        with await client.websocket_connect("/") as ws:
+            ws.send_text("Hello")
+            data_1 = ws.receive_text()
+            data_2 = ws.receive_text()
+            assert sorted([data_1, data_2]) == sorted(["Hello", "ping"])
+
+
+async def test_websocket_handler():
+    async with AsyncTestClient(app_stream_and_receive_raw) as client:
+        with await client.websocket_connect("/") as ws:
+            echo_data = {"data": "I should be in response"}
+            ws.send_json(echo_data)
+            assert ws.receive_json(timeout=0.5) == {"handle_receive": "start"}
+            assert ws.receive_json(timeout=0.5) == echo_data
+            assert ws.receive_text(timeout=0.5)