fix test run error

Author: hafid-Qa

src/fastapi_cli/cli.py

@@ -1,3 +1,5 @@
+import asyncio
+from enum import Enum
 from logging import getLogger
 from pathlib import Path
 from typing import Any, Union
@@ -16,6 +18,14 @@
 
 app = typer.Typer(rich_markup_mode="rich")
 
+
+class WSProtocolType(str, Enum):
+    auto = "auto"
+    none = "none"
+    websockets = "websockets"
+    wsproto = "wsproto"
+
+
 setup_logging()
 logger = getLogger(__name__)
 
@@ -60,6 +70,12 @@ def _run(
     command: str,
     app: Union[str, None] = None,
     proxy_headers: bool = False,
+    ws: type[asyncio.Protocol] | WSProtocolType = "auto",
+    ws_max_size: int = 16777216,
+    ws_max_queue: int = 32,
+    ws_ping_interval: float = 20.0,
+    ws_ping_timeout: float = 20.0,
+    ws_per_message_deflate: bool = True,
 ) -> None:
     try:
         use_uvicorn_app = get_import_string(path=path, app_name=app)
@@ -97,6 +113,12 @@ def _run(
         workers=workers,
         root_path=root_path,
         proxy_headers=proxy_headers,
+        ws=ws,
+        ws_max_size=ws_max_size,
+        ws_max_queue=ws_max_queue,
+        ws_ping_interval=ws_ping_interval,
+        ws_ping_timeout=ws_ping_timeout,
+        ws_per_message_deflate=ws_per_message_deflate,
     )
 
 
@@ -145,6 +167,32 @@ def dev(
             help="Enable/Disable X-Forwarded-Proto, X-Forwarded-For, X-Forwarded-Port to populate remote address info."
         ),
     ] = True,
+    ws: Annotated[
+        WSProtocolType,
+        typer.Option(
+            help="The WebSocket protocol.", case_sensitive=False, show_choices=True
+        ),
+    ] = WSProtocolType.auto,
+    ws_max_size: Annotated[
+        int,
+        typer.Option(help="WebSocket max size message in bytes."),
+    ] = 16777216,
+    ws_max_queue: Annotated[
+        int,
+        typer.Option(help="The maximum length of the WebSocket message queue."),
+    ] = 32,
+    ws_ping_interval: Annotated[
+        float,
+        typer.Option(help="WebSocket ping interval in seconds."),
+    ] = 20.0,
+    ws_ping_timeout: Annotated[
+        float,
+        typer.Option(help="WebSocket ping timeout in seconds."),
+    ] = 20.0,
+    ws_per_message_deflate: Annotated[
+        bool,
+        typer.Option(help="WebSocket per-message-deflate compression"),
+    ] = True,
 ) -> Any:
     """
     Run a [bold]FastAPI[/bold] app in [yellow]development[/yellow] mode. 🧪
@@ -180,6 +228,12 @@ def dev(
         app=app,
         command="dev",
         proxy_headers=proxy_headers,
+        ws=ws,
+        ws_max_size=ws_max_size,
+        ws_max_queue=ws_max_queue,
+        ws_ping_interval=ws_ping_interval,
+        ws_ping_timeout=ws_ping_timeout,
+        ws_per_message_deflate=ws_per_message_deflate,
     )
 
 
@@ -234,6 +288,32 @@ def run(
             help="Enable/Disable X-Forwarded-Proto, X-Forwarded-For, X-Forwarded-Port to populate remote address info."
         ),
     ] = True,
+    ws: Annotated[
+        WSProtocolType,
+        typer.Option(
+            help="The WebSocket protocol.", case_sensitive=False, show_choices=True
+        ),
+    ] = WSProtocolType.auto,
+    ws_max_size: Annotated[
+        int,
+        typer.Option(help="WebSocket max size message in bytes."),
+    ] = 16777216,
+    ws_max_queue: Annotated[
+        int,
+        typer.Option(help="The maximum length of the WebSocket message queue."),
+    ] = 32,
+    ws_ping_interval: Annotated[
+        float,
+        typer.Option(help="WebSocket ping interval in seconds."),
+    ] = 20.0,
+    ws_ping_timeout: Annotated[
+        float,
+        typer.Option(help="WebSocket ping timeout in seconds."),
+    ] = 20.0,
+    ws_per_message_deflate: Annotated[
+        bool,
+        typer.Option(help="WebSocket per-message-deflate compression"),
+    ] = True,
 ) -> Any:
     """
     Run a [bold]FastAPI[/bold] app in [green]production[/green] mode. 🚀
@@ -270,6 +350,12 @@ def run(
         app=app,
         command="run",
         proxy_headers=proxy_headers,
+        ws=ws,
+        ws_max_size=ws_max_size,
+        ws_max_queue=ws_max_queue,
+        ws_ping_interval=ws_ping_interval,
+        ws_ping_timeout=ws_ping_timeout,
+        ws_per_message_deflate=ws_per_message_deflate,
     )
 
 

tests/test_cli.py

@@ -29,6 +29,12 @@ def test_dev() -> None:
                 "workers": None,
                 "root_path": "",
                 "proxy_headers": True,
+                "ws": "auto",
+                "ws_max_size": 16777216,
+                "ws_max_queue": 32,
+                "ws_ping_interval": 20.0,
+                "ws_ping_timeout": 20.0,
+                "ws_per_message_deflate": True,
             }
         assert "Using import string single_file_app:app" in result.output
         assert (
@@ -58,6 +64,8 @@ def test_dev_args() -> None:
                     "--app",
                     "api",
                     "--no-proxy-headers",
+                    "--ws",
+                    "auto",
                 ],
             )
             assert result.exit_code == 0, result.output
@@ -71,6 +79,12 @@ def test_dev_args() -> None:
                 "workers": None,
                 "root_path": "/api",
                 "proxy_headers": False,
+                "ws": "auto",
+                "ws_max_size": 16777216,
+                "ws_max_queue": 32,
+                "ws_ping_interval": 20.0,
+                "ws_ping_timeout": 20.0,
+                "ws_per_message_deflate": True,
             }
         assert "Using import string single_file_app:api" in result.output
         assert (
@@ -97,6 +111,12 @@ def test_run() -> None:
                 "workers": None,
                 "root_path": "",
                 "proxy_headers": True,
+                "ws": "auto",
+                "ws_max_size": 16777216,
+                "ws_max_queue": 32,
+                "ws_ping_interval": 20.0,
+                "ws_ping_timeout": 20.0,
+                "ws_per_message_deflate": True,
             }
         assert "Using import string single_file_app:app" in result.output
         assert (
@@ -128,6 +148,8 @@ def test_run_args() -> None:
                     "--app",
                     "api",
                     "--no-proxy-headers",
+                    "--ws",
+                    "websockets",
                 ],
             )
             assert result.exit_code == 0, result.output
@@ -141,6 +163,12 @@ def test_run_args() -> None:
                 "workers": 2,
                 "root_path": "/api",
                 "proxy_headers": False,
+                "ws": "websockets",
+                "ws_max_size": 16777216,
+                "ws_max_queue": 32,
+                "ws_ping_interval": 20.0,
+                "ws_ping_timeout": 20.0,
+                "ws_per_message_deflate": True,
             }
         assert "Using import string single_file_app:api" in result.output
         assert (