Repository: encode/uvicorn

Issue:
Passing `reload_dirs` as a string instead of a list
@Roang-zero1 @euri10 @Kludex this PR introduced a regression and deleted the test for it (`test_reload_dir_is_set()`). Previously you could pass `reload_dirs` a string instead of a list. Now:
```python
Config('myapp:app', reload=True, reload_dirs='src')
```
prints:
```
WARNING:  Provided reload directories ['r', 'c', 's'] did not contain valid directories, watching current working directory.
```
https://github.com/encode/uvicorn/pull/820/files#diff-6982a5e270879281f9182303a9da9776148857b46325afb230e50ef8eb3264c9R294
(Of course, `set` expects an iterable... and strings are iterables of chars.)

_Originally posted by @HansBrende in https://github.com/encode/uvicorn/issues/820#issuecomment-977957988_

Discussion:

Patch:
diff --git a/uvicorn/config.py b/uvicorn/config.py
index 683a830..59c94b4 100644
--- a/uvicorn/config.py
+++ b/uvicorn/config.py
@@ -180,7 +180,15 @@ def resolve_reload_patterns(
 
     directories = list(set(directories).difference(set(children)))
 
-    return (list(set(patterns)), directories)
+    return list(set(patterns)), directories
+
+
+def _normalize_dirs(dirs: Union[List[str], str, None]) -> List[str]:
+    if dirs is None:
+        return []
+    if isinstance(dirs, str):
+        return [dirs]
+    return list(set(dirs))
 
 
 class Config:
@@ -291,9 +299,9 @@ class Config:
             )
 
         if self.should_reload:
-            reload_dirs = list(set(reload_dirs)) if reload_dirs else []
-            reload_includes = list(set(reload_includes)) if reload_includes else []
-            reload_excludes = list(set(reload_excludes)) if reload_excludes else []
+            reload_dirs = _normalize_dirs(reload_dirs)
+            reload_includes = _normalize_dirs(reload_includes)
+            reload_excludes = _normalize_dirs(reload_excludes)
 
             self.reload_includes, self.reload_dirs = resolve_reload_patterns(
                 reload_includes, reload_dirs
diff --git a/uvicorn/protocols/websockets/websockets_impl.py b/uvicorn/protocols/websockets/websockets_impl.py
index f4fa67f..7caabad 100644
--- a/uvicorn/protocols/websockets/websockets_impl.py
+++ b/uvicorn/protocols/websockets/websockets_impl.py
@@ -282,10 +282,17 @@ class WebSocketProtocol(_LoggerMixin, websockets.WebSocketServerProtocol):
             return {"type": "websocket.connect"}
 
         await self.handshake_completed_event.wait()
+
+        if self.closed_event.is_set():
+            # If client disconnected, use WebSocketServerProtocol.close_code property.
+            # If the handshake failed or the app closed before handshake completion,
+            # use 1006 Abnormal Closure.
+            return {"type": "websocket.disconnect", "code": self.close_code or 1006}
+
         try:
-            await self.ensure_open()
             data = await self.recv()
         except websockets.ConnectionClosed as exc:
+            self.closed_event.set()
             return {"type": "websocket.disconnect", "code": exc.code}
 
         msg = {"type": "websocket.receive"}


Test patch:
diff --git a/tests/protocols/test_websocket.py b/tests/protocols/test_websocket.py
index 2bd45b7..b444b3f 100644
--- a/tests/protocols/test_websocket.py
+++ b/tests/protocols/test_websocket.py
@@ -1,3 +1,5 @@
+import asyncio
+
 import httpx
 import pytest
 
@@ -543,3 +545,67 @@ async def test_send_binary_data_to_server_bigger_than_default(
             with pytest.raises(websockets.ConnectionClosedError) as e:
                 data = await send_text("ws://127.0.0.1:8000")
             assert e.value.code == expected_result
+
+
+@pytest.mark.asyncio
+@pytest.mark.parametrize("ws_protocol_cls", WS_PROTOCOLS)
+@pytest.mark.parametrize("http_protocol_cls", HTTP_PROTOCOLS)
+async def test_server_reject_connection(ws_protocol_cls, http_protocol_cls):
+    async def app(scope, receive, send):
+        assert scope["type"] == "websocket"
+
+        # Pull up first recv message.
+        message = await receive()
+        assert message["type"] == "websocket.connect"
+
+        # Reject the connection.
+        await send({"type": "websocket.close"})
+        # -- At this point websockets' recv() is unusable. --
+
+        # This doesn't raise `TypeError`:
+        # See https://github.com/encode/uvicorn/issues/244
+        message = await receive()
+        assert message["type"] == "websocket.disconnect"
+
+    async def websocket_session(url):
+        try:
+            async with websockets.connect(url):
+                pass
+        except Exception:
+            pass
+
+    config = Config(app=app, ws=ws_protocol_cls, http=http_protocol_cls, lifespan="off")
+    async with run_server(config):
+        await websocket_session("ws://127.0.0.1:8000")
+
+
+@pytest.mark.asyncio
+@pytest.mark.parametrize("ws_protocol_cls", WS_PROTOCOLS)
+@pytest.mark.parametrize("http_protocol_cls", HTTP_PROTOCOLS)
+async def test_server_can_read_messages_in_buffer_after_close(
+    ws_protocol_cls, http_protocol_cls
+):
+    frames = []
+
+    class App(WebSocketResponse):
+        async def websocket_connect(self, message):
+            await self.send({"type": "websocket.accept"})
+            # Ensure server doesn't start reading frames from read buffer until
+            # after client has sent close frame, but server is still able to
+            # read these frames
+            await asyncio.sleep(0.2)
+
+        async def websocket_receive(self, message):
+            frames.append(message.get("bytes"))
+
+    async def send_text(url):
+        async with websockets.connect(url) as websocket:
+            await websocket.send(b"abc")
+            await websocket.send(b"abc")
+            await websocket.send(b"abc")
+
+    config = Config(app=App, ws=ws_protocol_cls, http=http_protocol_cls, lifespan="off")
+    async with run_server(config):
+        await send_text("ws://127.0.0.1:8000")
+
+    assert frames == [b"abc", b"abc", b"abc"]
diff --git a/tests/test_config.py b/tests/test_config.py
index 7a04092..265d09b 100644
--- a/tests/test_config.py
+++ b/tests/test_config.py
@@ -119,6 +119,10 @@ def test_reload_dir_is_set(
             == f"Will watch for changes in these directories: {[str(app_dir)]}"
         )
         assert config.reload_dirs == [app_dir]
+        config = Config(
+            app="tests.test_config:asgi_app", reload=True, reload_dirs=str(app_dir)
+        )
+        assert config.reload_dirs == [app_dir]
 
 
 def test_non_existant_reload_dir_is_not_set(