test: using anyio instead of asyncio in tests

WSH032 · WSH032 · commit ef8da551f4af · 2024-04-03T06:46:46.000+08:00
And simplify the `UvicornServer` class in test
diff --git a/pyproject.toml b/pyproject.toml
@@ -98,9 +98,11 @@ dependencies = [
     "pytest == 7.*",
     "pytest-cov ==  4.*",
     "uvicorn[standard] < 1.0.0", # TODO: Once it releases version 1.0.0, we will remove this restriction.
+    "hypercorn[trio] == 0.16.*",
     "httpx[http2]",              # we don't set version here, instead set it in `[project].dependencies`.
-    "asgi-lifespan==2.*",
-    "pytest-timeout==2.*",
+    "asgi-lifespan == 2.*",
+    "pytest-timeout == 2.*",
+    "sniffio == 1.3.*",
 ]
 
 [tool.hatch.envs.default.scripts]
diff --git a/tests/app/echo_ws_app.py b/tests/app/echo_ws_app.py
@@ -1,8 +1,8 @@
 # ruff: noqa: D100
 # pyright: reportUnusedFunction=false
 
-import asyncio
 
+import anyio
 from fastapi import FastAPI, WebSocket
 from starlette.websockets import WebSocketDisconnect
 
@@ -76,7 +76,7 @@ async def just_close_with_1001(websocket: WebSocket):
         test_app_dataclass.request_dict["request"] = websocket
 
         await websocket.accept()
-        await asyncio.sleep(0.3)
+        await anyio.sleep(0.3)
         await websocket.close(1001)
 
     @app.websocket("/reject_handshake")
diff --git a/tests/app/tool.py b/tests/app/tool.py
@@ -1,18 +1,16 @@
 # noqa: D100
 
-import asyncio
-import socket
+from contextlib import AsyncExitStack
 from dataclasses import dataclass
-from typing import Any, Callable, List, Optional, Type, TypedDict, TypeVar, Union
+from typing import Any, TypedDict, Union
 
+import anyio
 import httpx
 import uvicorn
 from fastapi import FastAPI
 from starlette.requests import Request
 from starlette.websockets import WebSocket
-from typing_extensions import Self, override
-
-_Decoratable_T = TypeVar("_Decoratable_T", bound=Union[Callable[..., Any], Type[Any]])
+from typing_extensions import Self
 
 ServerRecvRequestsTypes = Union[Request, WebSocket]
 
@@ -46,180 +44,32 @@ def get_request(self) -> ServerRecvRequestsTypes:
         return server_recv_request
 
 
-def _no_override_uvicorn_server(_method: _Decoratable_T) -> _Decoratable_T:
-    """Check if the method is already in `uvicorn.Server`."""
-    assert not hasattr(
-        uvicorn.Server, _method.__name__
-    ), f"Override method of `uvicorn.Server` cls : {_method.__name__}"
-    return _method
-
-
-class AeixtTimeoutUndefine:
-    """Didn't set `contx_exit_timeout` in `aexit()`."""
-
-
-aexit_timeout_undefine = AeixtTimeoutUndefine()
-
-
-# HACK: 不能继承 AbstractAsyncContextManager[Self]
-# 目前有问题，继承 AbstractAsyncContextManager 的话pyright也推测不出来类型
-# 只能依靠 __aenter__ 和 __aexit__ 的类型注解
 class UvicornServer(uvicorn.Server):
-    """subclass of `uvicorn.Server` which can use AsyncContext to launch and shutdown automatically.
-
-    Attributes:
-        contx_server_task: The task of server.
-        contx_socket: The socket of server.
-
-        other attributes are same as `uvicorn.Server`:
-            - config: The config arg that be passed in.
-            ...
-    """
+    """subclass of `uvicorn.Server` which can use AsyncContext to launch and shutdown automatically."""
 
-    _contx_server_task: Union["asyncio.Task[None]", None]
-    assert not hasattr(uvicorn.Server, "_contx_server_task")
-
-    _contx_socket: Union[socket.socket, None]
-    assert not hasattr(uvicorn.Server, "_contx_socket")
-
-    _contx_server_started_event: Union[asyncio.Event, None]
-    assert not hasattr(uvicorn.Server, "_contx_server_started_event")
-
-    contx_exit_timeout: Union[int, float, None]
-    assert not hasattr(uvicorn.Server, "contx_exit_timeout")
-
-    @override
-    def __init__(
-        self, config: uvicorn.Config, contx_exit_timeout: Union[int, float, None] = None
-    ) -> None:
-        """The same as `uvicorn.Server.__init__`."""
-        super().__init__(config=config)
-        self._contx_server_task = None
-        self._contx_socket = None
-        self._contx_server_started_event = None
-        self.contx_exit_timeout = contx_exit_timeout
-
-    @override
-    async def startup(self, sockets: Optional[List[socket.socket]] = None) -> None:
-        """The same as `uvicorn.Server.startup`."""
-        super_return = await super().startup(sockets=sockets)
-        self.contx_server_started_event.set()
-        return super_return
-
-    @_no_override_uvicorn_server
-    async def aenter(self) -> Self:
+    async def __aenter__(self) -> Self:
         """Launch the server."""
-        # 在分配资源之前，先检查是否重入
-        if self.contx_server_started_event.is_set():
-            raise RuntimeError("DO not launch server by __aenter__ again!")
-
         # FIXME: # 这个socket被设计为可被同一进程内的多个server共享，可能会引起潜在问题
-        self._contx_socket = self.config.bind_socket()
+        self._socket = self.config.bind_socket()
+        self._exit_stack = AsyncExitStack()
 
-        self._contx_server_task = asyncio.create_task(
-            self.serve([self._contx_socket]), name=f"Uvicorn Server Task of {self}"
+        task_group = await self._exit_stack.enter_async_context(
+            anyio.create_task_group()
+        )
+        task_group.start_soon(
+            self.serve, [self._socket], name=f"Uvicorn Server Task of {self}"
         )
-        # 在 uvicorn.Server 的实现中，Server.serve() 内部会调用 Server.startup() 完成启动
-        # 被覆盖的 self.startup() 会在完成时调用 self.contx_server_started_event.set()
-        await self.contx_server_started_event.wait()  # 等待服务器确实启动后才返回
-        return self
 
-    @_no_override_uvicorn_server
-    async def __aenter__(self) -> Self:
-        """Launch the server.
+        return self
 
-        The same as `self.aenter()`.
-        """
-        return await self.aenter()
-
-    @_no_override_uvicorn_server
-    async def aexit(
-        self,
-        contx_exit_timeout: Union[
-            int, float, None, AeixtTimeoutUndefine
-        ] = aexit_timeout_undefine,
-    ) -> None:
+    async def __aexit__(self, *_: Any, **__: Any) -> None:
         """Shutdown the server."""
-        contx_server_task = self.contx_server_task
-        contx_socket = self.contx_socket
-
-        if isinstance(contx_exit_timeout, AeixtTimeoutUndefine):
-            contx_exit_timeout = self.contx_exit_timeout
-
         # 在 uvicorn.Server 的实现中，设置 should_exit 可以使得 server 任务结束
-        assert hasattr(self, "should_exit")
+        assert self.should_exit is False, "The server has already exited."
         self.should_exit = True
-
-        try:
-            await asyncio.wait_for(contx_server_task, timeout=contx_exit_timeout)
-        except asyncio.TimeoutError:
-            print(f"{contx_server_task.get_name()} timeout!")
-        finally:
-            # 其实uvicorn.Server会自动关闭socket，这里是为了保险起见
-            contx_socket.close()
-
-    @_no_override_uvicorn_server
-    async def __aexit__(self, *_: Any, **__: Any) -> None:
-        """Shutdown the server.
-
-        The same as `self.aexit()`.
-        """
-        return await self.aexit()
-
-    @property
-    @_no_override_uvicorn_server
-    def contx_server_started_event(self) -> asyncio.Event:
-        """The event that indicates the server has started.
-
-        When first call the property, it will instantiate a `asyncio.Event()`to
-        `self._contx_server_started_event`.
-
-        Warn: This is a internal implementation detail, do not change the event manually.
-            - please call the property in `self.aenter()` or `self.startup()` **first**.
-            - **Never** call it outside of an async event loop first:
-                https://stackoverflow.com/questions/53724665/using-queues-results-in-asyncio-exception-got-future-future-pending-attached
-        """
-        if self._contx_server_started_event is None:
-            self._contx_server_started_event = asyncio.Event()
-
-        return self._contx_server_started_event
-
-    @property
-    @_no_override_uvicorn_server
-    def contx_socket(self) -> socket.socket:
-        """The socket of server.
-
-        Note: must call `self.__aenter__()` first.
-        """
-        if self._contx_socket is None:
-            raise RuntimeError("Please call `self.__aenter__()` first.")
-        else:
-            return self._contx_socket
-
-    @property
-    @_no_override_uvicorn_server
-    def contx_server_task(self) -> "asyncio.Task[None]":
-        """The task of server.
-
-        Note: must call `self.__aenter__()` first.
-        """
-        if self._contx_server_task is None:
-            raise RuntimeError("Please call `self.__aenter__()` first.")
-        else:
-            return self._contx_server_task
-
-    @property
-    @_no_override_uvicorn_server
-    def contx_socket_getname(self) -> Any:
-        """Utils for calling self.contx_socket.getsockname().
-
-        Return:
-            refer to: https://docs.python.org/zh-cn/3/library/socket.html#socket-families
-        """
-        return self.contx_socket.getsockname()
+        await self._exit_stack.__aexit__(*_, **__)
 
     @property
-    @_no_override_uvicorn_server
     def contx_socket_url(self) -> httpx.URL:
         """If server is tcp socket, return the url of server.
 
@@ -228,7 +78,8 @@ def contx_socket_url(self) -> httpx.URL:
         config = self.config
         if config.fd is not None or config.uds is not None:
             raise RuntimeError("Only support tcp socket.")
-        host, port = self.contx_socket_getname[:2]
+        # refer to: https://docs.python.org/zh-cn/3/library/socket.html#socket-families
+        host, port = self._socket.getsockname()[:2]
         return httpx.URL(
             host=host,
             port=port,
diff --git a/tests/conftest.py b/tests/conftest.py
@@ -16,7 +16,6 @@
     Coroutine,
     Literal,
     Protocol,
-    Union,
 )
 
 import pytest
@@ -64,7 +63,7 @@ class LifeAppDataclass4Test(AppDataclass4Test):
 
 class UvicornServerFixture(Protocol):  # noqa: D101
     def __call__(  # noqa: D102
-        self, config: uvicorn.Config, contx_exit_timeout: Union[int, float, None] = None
+        self, config: uvicorn.Config
     ) -> Coroutine[None, None, UvicornServer]: ...
 
 
@@ -199,11 +198,9 @@ async def uvicorn_server_fixture() -> AsyncIterator[UvicornServerFixture]:
     """
     async with AsyncExitStack() as exit_stack:
 
-        async def uvicorn_server_fct(
-            config: uvicorn.Config, contx_exit_timeout: Union[int, float, None] = None
-        ) -> UvicornServer:
+        async def uvicorn_server_fct(config: uvicorn.Config) -> UvicornServer:
             uvicorn_server = await exit_stack.enter_async_context(
-                UvicornServer(config=config, contx_exit_timeout=contx_exit_timeout)
+                UvicornServer(config=config)
             )
             return uvicorn_server
 
diff --git a/tests/test_ws.py b/tests/test_ws.py
@@ -1,11 +1,11 @@
 # noqa: D100
 
 
-import asyncio
 from contextlib import AsyncExitStack
 from multiprocessing import Process, Queue
 from typing import Any, Dict, Literal, Optional
 
+import anyio
 import httpx
 import httpx_ws
 import pytest
@@ -25,7 +25,6 @@
 
 DEFAULT_HOST = "127.0.0.1"
 DEFAULT_PORT = 0
-DEFAULT_CONTX_EXIT_TIMEOUT = 5
 
 # WS_BACKENDS_NEED_BE_TESTED = ("websockets", "wsproto")
 # # FIXME: wsproto 有问题，暂时不测试
@@ -56,14 +55,14 @@ def _subprocess_run_echo_ws_uvicorn_server(queue: "Queue[str]", **kwargs: Any):
     )
 
     async def run():
-        await target_ws_server.aenter()
-        url = str(target_ws_server.contx_socket_url)
-        queue.put(url)
-        queue.close()
-        while True:  # run forever
-            await asyncio.sleep(0.1)
+        async with target_ws_server:
+            url = str(target_ws_server.contx_socket_url)
+            queue.put(url)
+            queue.close()
+            while True:  # run forever
+                await anyio.sleep(0.1)
 
-    asyncio.run(run())
+    anyio.run(run)
 
 
 def _subprocess_run_httpx_ws(
@@ -96,9 +95,9 @@ async def run():
         queue.put("done")
         queue.close()
         while True:  # run forever
-            await asyncio.sleep(0.1)
+            await anyio.sleep(0.1)
 
-    asyncio.run(run())
+    anyio.run(run)
 
 
 class TestReverseWsProxy(AbstractTestProxy):
@@ -120,7 +119,6 @@ async def tool_4_test_fixture(  # pyright: ignore[reportIncompatibleMethodOverri
             uvicorn.Config(
                 echo_ws_app, port=DEFAULT_PORT, host=DEFAULT_HOST, ws=request.param
             ),
-            contx_exit_timeout=DEFAULT_CONTX_EXIT_TIMEOUT,
         )
 
         target_server_base_url = str(target_ws_server.contx_socket_url)
@@ -135,7 +133,6 @@ async def tool_4_test_fixture(  # pyright: ignore[reportIncompatibleMethodOverri
             uvicorn.Config(
                 reverse_ws_app, port=DEFAULT_PORT, host=DEFAULT_HOST, ws=request.param
             ),
-            contx_exit_timeout=DEFAULT_CONTX_EXIT_TIMEOUT,
         )
 
         proxy_server_base_url = str(proxy_ws_server.contx_socket_url)
@@ -226,7 +223,7 @@ async def test_ws_proxy(self, tool_4_test_fixture: Tool4TestFixture) -> None:
 
         # 避免从队列中get导致的异步阻塞
         while aconnect_ws_subprocess_queue.empty():
-            await asyncio.sleep(0.1)
+            await anyio.sleep(0.1)
         _ = aconnect_ws_subprocess_queue.get()  # 获取到了即代表连接建立成功
 
         # force shutdown client
@@ -267,7 +264,7 @@ async def test_target_server_shutdown_abnormally(
 
         # 避免从队列中get导致的异步阻塞
         while subprocess_queue.empty():
-            await asyncio.sleep(0.1)
+            await anyio.sleep(0.1)
         target_server_base_url = subprocess_queue.get()
 
         client_for_conn_to_target_server = httpx.AsyncClient(proxies=NO_PROXIES)
@@ -300,13 +297,11 @@ async def test_target_server_shutdown_abnormally(
                     await ws0.receive()
                 assert exce.value.code == 1011
 
-                loop = asyncio.get_running_loop()
-
-                seconde_ws_recv_start = loop.time()
+                seconde_ws_recv_start = anyio.current_time()
                 with pytest.raises(httpx_ws.WebSocketDisconnect) as exce:
                     await ws1.receive()
                 assert exce.value.code == 1011
-                seconde_ws_recv_end = loop.time()
+                seconde_ws_recv_end = anyio.current_time()
 
                 # HACK: 由于收到关闭代码需要40s，目前无法确定是什么原因，
                 # 所以目前会同时测试两个客户端的连接，
