Repository: litestar-org/litestar

Issue:
Bug: Arbitrary dependencies in `WebSocketListener` fail
### Description

Injecting arbitrary dependencies (that is, dependencies that are not "reserved" like `state`) cannot be successfully injected into a `WebSocketListener` subclass. The injection works as intended when using the `websocket_listener` decorator.


### MCVE

```python
from litestar import WebSocket, Litestar
from litestar.handlers import WebsocketListener
from litestar.testing import TestClient


class Listener(WebsocketListener):
    path = "/ws"

    def on_accept(self, socket: WebSocket, foo: str) -> None:
        pass

    async def on_receive(self, data: dict) -> dict:
        return data


app = Litestar(
    [Listener],
    dependencies={"foo": lambda: "bar"},
)


with TestClient(app) as client, client.websocket_connect("/ws") as ws:
    ws.send_json({"foo": "bar"})
```


### Litestar Version

2.0.0alpha7
Bug: Arbitrary dependencies in `WebSocketListener` fail
### Description

Injecting arbitrary dependencies (that is, dependencies that are not "reserved" like `state`) cannot be successfully injected into a `WebSocketListener` subclass. The injection works as intended when using the `websocket_listener` decorator.


### MCVE

```python
from litestar import WebSocket, Litestar
from litestar.handlers import WebsocketListener
from litestar.testing import TestClient


class Listener(WebsocketListener):
    path = "/ws"

    def on_accept(self, socket: WebSocket, foo: str) -> None:
        pass

    async def on_receive(self, data: dict) -> dict:
        return data


app = Litestar(
    [Listener],
    dependencies={"foo": lambda: "bar"},
)


with TestClient(app) as client, client.websocket_connect("/ws") as ws:
    ws.send_json({"foo": "bar"})
```


### Litestar Version

2.0.0alpha7


Discussion:

Resolving patch:
diff --git a/litestar/handlers/websocket_handlers/listener.py b/litestar/handlers/websocket_handlers/listener.py
--- a/litestar/handlers/websocket_handlers/listener.py
+++ b/litestar/handlers/websocket_handlers/listener.py
@@ -42,7 +42,7 @@
 if TYPE_CHECKING:
     from typing import Coroutine
 
-    from litestar import Litestar
+    from litestar import Litestar, Router
     from litestar.dto.interface import DTOInterface
     from litestar.types.asgi_types import WebSocketMode
 
@@ -358,8 +358,16 @@ class WebsocketListener(ABC):
     type_encoders: A mapping of types to callables that transform them into types supported for serialization.
     """
 
-    def __init__(self) -> None:
-        self._handler = websocket_listener(
+    def __init__(self, owner: Router) -> None:
+        """Initialize a WebsocketListener instance.
+
+        Args:
+            owner: The :class:`Router <.router.Router>` instance that owns this listener.
+        """
+        self._owner = owner
+
+    def to_handler(self) -> websocket_listener:
+        handler = websocket_listener(
             dependencies=self.dependencies,
             dto=self.dto,
             exception_handlers=self.exception_handlers,
@@ -376,6 +384,8 @@ def __init__(self) -> None:
             signature_namespace=self.signature_namespace,
             type_encoders=self.type_encoders,
         )(self.on_receive)
+        handler.owner = self._owner
+        return handler
 
     @abstractmethod
     def on_receive(self, *args: Any, **kwargs: Any) -> Any:
diff --git a/litestar/router.py b/litestar/router.py
--- a/litestar/router.py
+++ b/litestar/router.py
@@ -290,7 +290,7 @@ def _validate_registration_value(self, value: ControllerRouterHandler) -> Contro
 
         # this narrows down to an ABC, but we assume a non-abstract subclass of the ABC superclass
         if is_class_and_subclass(value, WebsocketListener):  # type: ignore[type-abstract]
-            return value()._handler  # pyright: ignore
+            return value(owner=self).to_handler()  # pyright: ignore
 
         if isinstance(value, Router):
             if value.owner:


Test patch:
diff --git a/tests/handlers/websocket/test_listeners.py b/tests/handlers/websocket/test_listeners.py
--- a/tests/handlers/websocket/test_listeners.py
+++ b/tests/handlers/websocket/test_listeners.py
@@ -364,25 +364,72 @@ def test_hook_dependencies() -> None:
     on_accept_mock = MagicMock()
     on_disconnect_mock = MagicMock()
 
-    def on_accept(name: str, state: State, query: dict) -> None:
-        on_accept_mock(name=name, state=state, query=query)
+    def some_dependency() -> str:
+        return "hello"
 
-    def on_disconnect(name: str, state: State, query: dict) -> None:
-        on_disconnect_mock(name=name, state=state, query=query)
+    def on_accept(name: str, state: State, query: dict, some: str) -> None:
+        on_accept_mock(name=name, state=state, query=query, some=some)
+
+    def on_disconnect(name: str, state: State, query: dict, some: str) -> None:
+        on_disconnect_mock(name=name, state=state, query=query, some=some)
 
     @websocket_listener("/{name: str}", on_accept=on_accept, on_disconnect=on_disconnect)
     def handler(data: bytes) -> None:
         pass
 
-    with create_test_client([handler], debug=True) as client, client.websocket_connect("/foo") as ws:
+    with create_test_client(
+        [handler], debug=True, dependencies={"some": some_dependency}
+    ) as client, client.websocket_connect("/foo") as ws:
+        ws.send_text("")
+
+    on_accept_kwargs = on_accept_mock.call_args_list[0].kwargs
+    assert on_accept_kwargs["name"] == "foo"
+    assert on_accept_kwargs["some"] == "hello"
+    assert isinstance(on_accept_kwargs["state"], State)
+    assert isinstance(on_accept_kwargs["query"], dict)
+
+    on_disconnect_kwargs = on_disconnect_mock.call_args_list[0].kwargs
+    assert on_disconnect_kwargs["name"] == "foo"
+    assert on_disconnect_kwargs["some"] == "hello"
+    assert isinstance(on_disconnect_kwargs["state"], State)
+    assert isinstance(on_disconnect_kwargs["query"], dict)
+
+
+def test_websocket_listener_class_hook_dependencies() -> None:
+    on_accept_mock = MagicMock()
+    on_disconnect_mock = MagicMock()
+
+    def some_dependency() -> str:
+        return "hello"
+
+    class Listener(WebsocketListener):
+        path = "/{name: str}"
+
+        def on_accept(self, name: str, state: State, query: dict, some: str) -> None:  # type: ignore[override]
+            on_accept_mock(name=name, state=state, query=query, some=some)
+
+        def on_disconnect(self, name: str, state: State, query: dict, some: str) -> None:  # type: ignore[override]
+            on_disconnect_mock(name=name, state=state, query=query, some=some)
+
+        def on_receive(self, data: bytes) -> None:
+            pass
+
+    with create_test_client(
+        [Listener], debug=True, dependencies={"some": some_dependency}
+    ) as client, client.websocket_connect("/foo") as ws:
         ws.send_text("")
 
-    assert on_accept_mock.call_args_list[0].kwargs["name"] == "foo"
-    assert isinstance(on_accept_mock.call_args_list[0].kwargs["state"], State)
-    assert isinstance(on_accept_mock.call_args_list[0].kwargs["query"], dict)
-    assert on_disconnect_mock.call_args_list[0].kwargs["name"] == "foo"
-    assert isinstance(on_disconnect_mock.call_args_list[0].kwargs["state"], State)
-    assert isinstance(on_disconnect_mock.call_args_list[0].kwargs["query"], dict)
+    on_accept_kwargs = on_accept_mock.call_args_list[0].kwargs
+    assert on_accept_kwargs["name"] == "foo"
+    assert on_accept_kwargs["some"] == "hello"
+    assert isinstance(on_accept_kwargs["state"], State)
+    assert isinstance(on_accept_kwargs["query"], dict)
+
+    on_disconnect_kwargs = on_disconnect_mock.call_args_list[0].kwargs
+    assert on_disconnect_kwargs["name"] == "foo"
+    assert on_disconnect_kwargs["some"] == "hello"
+    assert isinstance(on_disconnect_kwargs["state"], State)
+    assert isinstance(on_disconnect_kwargs["query"], dict)
 
 
 @pytest.mark.parametrize("hook_name", ["on_accept", "on_disconnect", "connection_accept_handler"])