Repository: litestar-org/litestar

Issue:
Bug: Using `websocket_listener` in controller causes `TypeError: .handler_fn() got multiple values for argument 'socket'`
### Description

Confirmed by @provinzkraut in Discord



### URL to code causing the issue

_No response_

### MCVE

```python
from litestar import Controller, WebSocket, websocket_listener
from litestar.testing import create_test_client


class ClientController(Controller):
    path: str = "/client"

    @websocket_listener("/ws")
    async def websocket_handler(self, data: bytes, socket: WebSocket) -> bytes:
        return data


with create_test_client(ClientController) as client, client.websocket_connect("/client/ws") as ws:
    data = ws.send_bytes(b"foo")
    assert data == b"foo"
```


### Screenshots

"![Exception Log](https://user-images.githubusercontent.com/37921711/235967617-91018971-d13d-489f-b023-bf9fa807c3ae.png)"


### Logs

_No response_

### Litestar Version

2.0.0alpha5

### Platform

- [X] Linux
- [ ] Mac
- [ ] Windows
- [ ] Other (Please specify in the description above)

<!-- POLAR PLEDGE BADGE START -->
---
> [!NOTE]  
> While we are open for sponsoring on [GitHub Sponsors](https://github.com/sponsors/litestar-org/) and 
> [OpenCollective](https://opencollective.com/litestar), we also utilize [Polar.sh](https://polar.sh/) to engage in pledge-based sponsorship.
>
> Check out all issues funded or available for funding [on our Polar.sh dashboard](https://polar.sh/litestar-org)
> * If you would like to see an issue prioritized, make a pledge towards it!
> * We receive the pledge once the issue is completed & verified
> * This, along with engagement in the community, helps us know which features are a priority to our users.

<a href="https://polar.sh/litestar-org/litestar/issues/1615">
<picture>
  <source media="(prefers-color-scheme: dark)" srcset="https://polar.sh/api/github/litestar-org/litestar/issues/1615/pledge.svg?darkmode=1">
  <img alt="Fund with Polar" src="https://polar.sh/api/github/litestar-org/litestar/issues/1615/pledge.svg">
</picture>
</a>
<!-- POLAR PLEDGE BADGE END -->

StaticFilesConfig and virtual directories
I'm trying to write a ``FileSystemProtocol`` to load files from the package data using [importlib_resources](https://importlib-resources.readthedocs.io/en/latest/using.html#). But because ``directories`` is defined as ``DirectoryPath``, pydantic checks if the given directories exist in the local filesystem. 

This is not generally true, especially in any kind of virtual filesystem (e.g. a zipped package). I think this condition should be relaxed to support virtual filesystems.

https://github.com/starlite-api/starlite/blob/9bb6dcd57c10a591377cf8e3a537e9292566d5b9/starlite/config/static_files.py#L32


Discussion:

I agree. Do you want to submit a PR?

Resolving patch:
diff --git a/litestar/handlers/websocket_handlers/_utils.py b/litestar/handlers/websocket_handlers/_utils.py
--- a/litestar/handlers/websocket_handlers/_utils.py
+++ b/litestar/handlers/websocket_handlers/_utils.py
@@ -116,7 +116,7 @@ def create_handler_function(
 ) -> Callable[..., Coroutine[None, None, None]]:
     listener_callback = AsyncCallable(listener_context.listener_callback)
 
-    async def handler_fn(socket: WebSocket, **kwargs: Any) -> None:
+    async def handler_fn(*args: Any, socket: WebSocket, **kwargs: Any) -> None:
         ctx = ConnectionContext.from_connection(socket)
         data_dto = listener_context.resolved_data_dto(ctx) if listener_context.resolved_data_dto else None
         return_dto = listener_context.resolved_return_dto(ctx) if listener_context.resolved_return_dto else None
@@ -129,7 +129,7 @@ async def handler_fn(socket: WebSocket, **kwargs: Any) -> None:
         async with lifespan_manager(socket):
             while True:
                 received_data = await handle_receive(socket, data_dto)
-                data_to_send = await listener_callback(data=received_data, **kwargs)
+                data_to_send = await listener_callback(*args, data=received_data, **kwargs)
                 if handle_send:
                     await handle_send(socket, data_to_send, return_dto)
 
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
@@ -18,7 +18,7 @@
 from litestar._signature import create_signature_model
 from litestar.connection import WebSocket
 from litestar.dto.interface import HandlerContext
-from litestar.exceptions import ImproperlyConfiguredException
+from litestar.exceptions import ImproperlyConfiguredException, WebSocketDisconnect
 from litestar.serialization import default_serializer
 from litestar.types import (
     AnyCallable,
@@ -162,6 +162,8 @@ async def default_connection_lifespan(self, socket: WebSocket) -> AsyncGenerator
             await self.on_accept(socket)
         try:
             yield
+        except WebSocketDisconnect:
+            pass
         finally:
             if self.on_disconnect:
                 await self.on_disconnect(socket)


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
@@ -6,7 +6,7 @@
 import pytest
 from pytest_lazyfixture import lazy_fixture
 
-from litestar import Litestar, Request, WebSocket
+from litestar import Controller, Litestar, Request, WebSocket
 from litestar.datastructures import State
 from litestar.di import Provide
 from litestar.dto.factory import dto_field
@@ -316,3 +316,19 @@ def handler(data: bytes) -> None:
 
     on_accept.assert_called_once()
     on_disconnect.assert_called_once()
+
+
+def test_listener_in_controller() -> None:
+    # test for https://github.com/litestar-org/litestar/issues/1615
+
+    class ClientController(Controller):
+        path: str = "/"
+
+        @websocket_listener("/ws")
+        async def websocket_handler(self, data: str, socket: WebSocket) -> str:
+            return data
+
+    with create_test_client(ClientController, debug=True) as client, client.websocket_connect("/ws") as ws:
+        ws.send_text("foo")
+        data = ws.receive_text(timeout=1)
+        assert data == "foo"