Avoid asyncio dependency in tests

Author: praboud-ant

tests/server/fastmcp/auth/streaming_asgi_transport.py

@@ -6,12 +6,13 @@
 the connection is closed.
 """
 
-import asyncio
 import typing
 from typing import Any, Dict, Tuple
 
 import anyio
+import anyio.abc
 import anyio.streams.memory
+
 from httpx._models import Request, Response
 from httpx._transports.base import AsyncBaseTransport
 from httpx._types import AsyncByteStream
@@ -41,6 +42,7 @@ class StreamingASGITransport(AsyncBaseTransport):
     def __init__(
         self,
         app: typing.Callable,
+        task_group: anyio.abc.TaskGroup,
         raise_app_exceptions: bool = True,
         root_path: str = "",
         client: Tuple[str, int] = ("127.0.0.1", 123),
@@ -49,6 +51,7 @@ def __init__(
         self.raise_app_exceptions = raise_app_exceptions
         self.root_path = root_path
         self.client = client
+        self.task_group = task_group
 
     async def handle_async_request(
         self,
@@ -161,8 +164,8 @@ async def process_messages() -> None:
                 response_complete.set()
 
         # Create tasks for running the app and processing messages
-        asyncio.create_task(run_app())
-        asyncio.create_task(process_messages())
+        self.task_group.start_soon(run_app)
+        self.task_group.start_soon(process_messages)
 
         # Wait for the initial response or timeout
         await initial_response_ready.wait()

tests/server/fastmcp/auth/test_auth_integration.py

@@ -11,6 +11,7 @@
 from typing import List, Optional
 from urllib.parse import parse_qs, urlparse
 
+import anyio
 import httpx
 import pytest
 from httpx_sse import aconnect_sse
@@ -993,130 +994,132 @@ async def test_fastmcp_with_auth(
         def test_tool(x: int) -> str:
             return f"Result: {x}"
 
-        transport = StreamingASGITransport(app=mcp.starlette_app())  # pyright: ignore
-        test_client = httpx.AsyncClient(
-            transport=transport, base_url="http://mcptest.com"
-        )
-        # test_client = httpx.AsyncClient(app=mcp.starlette_app(), base_url="http://mcptest.com")
+        async with anyio.create_task_group() as task_group:
+            transport = StreamingASGITransport(app=mcp.starlette_app(), task_group=task_group)  # pyright: ignore
+            test_client = httpx.AsyncClient(
+                transport=transport, base_url="http://mcptest.com"
+            )
+            # test_client = httpx.AsyncClient(app=mcp.starlette_app(), base_url="http://mcptest.com")
 
-        # Test metadata endpoint
-        response = await test_client.get("/.well-known/oauth-authorization-server")
-        assert response.status_code == 200
+            # Test metadata endpoint
+            response = await test_client.get("/.well-known/oauth-authorization-server")
+            assert response.status_code == 200
 
-        # Test that auth is required for protected endpoints
-        response = await test_client.get("/sse")
-        # TODO: we should return 401/403 depending on whether authn or authz fails
-        assert response.status_code == 403
+            # Test that auth is required for protected endpoints
+            response = await test_client.get("/sse")
+            # TODO: we should return 401/403 depending on whether authn or authz fails
+            assert response.status_code == 403
 
-        response = await test_client.post("/messages/")
-        # TODO: we should return 401/403 depending on whether authn or authz fails
-        assert response.status_code == 403, response.content
+            response = await test_client.post("/messages/")
+            # TODO: we should return 401/403 depending on whether authn or authz fails
+            assert response.status_code == 403, response.content
 
-        response = await test_client.post(
-            "/messages/",
-            headers={"Authorization": "invalid"},
-        )
-        assert response.status_code == 403
-
-        response = await test_client.post(
-            "/messages/",
-            headers={"Authorization": "Bearer invalid"},
-        )
-        assert response.status_code == 403
+            response = await test_client.post(
+                "/messages/",
+                headers={"Authorization": "invalid"},
+            )
+            assert response.status_code == 403
 
-        # now, become authenticated and try to go through the flow again
-        client_metadata = {
-            "redirect_uris": ["https://client.example.com/callback"],
-            "client_name": "Test Client",
-        }
+            response = await test_client.post(
+                "/messages/",
+                headers={"Authorization": "Bearer invalid"},
+            )
+            assert response.status_code == 403
 
-        response = await test_client.post(
-            "/register",
-            json=client_metadata,
-        )
-        assert response.status_code == 201
-        client_info = response.json()
+            # now, become authenticated and try to go through the flow again
+            client_metadata = {
+                "redirect_uris": ["https://client.example.com/callback"],
+                "client_name": "Test Client",
+            }
 
-        # Request authorization using POST with form-encoded data
-        response = await test_client.post(
-            "/authorize",
-            data={
-                "response_type": "code",
-                "client_id": client_info["client_id"],
-                "redirect_uri": "https://client.example.com/callback",
-                "code_challenge": pkce_challenge["code_challenge"],
-                "code_challenge_method": "S256",
-                "state": "test_state",
-            },
-        )
-        assert response.status_code == 302
+            response = await test_client.post(
+                "/register",
+                json=client_metadata,
+            )
+            assert response.status_code == 201
+            client_info = response.json()
 
-        # Extract the authorization code from the redirect URL
-        redirect_url = response.headers["location"]
-        parsed_url = urlparse(redirect_url)
-        query_params = parse_qs(parsed_url.query)
+            # Request authorization using POST with form-encoded data
+            response = await test_client.post(
+                "/authorize",
+                data={
+                    "response_type": "code",
+                    "client_id": client_info["client_id"],
+                    "redirect_uri": "https://client.example.com/callback",
+                    "code_challenge": pkce_challenge["code_challenge"],
+                    "code_challenge_method": "S256",
+                    "state": "test_state",
+                },
+            )
+            assert response.status_code == 302
 
-        assert "code" in query_params
-        auth_code = query_params["code"][0]
+            # Extract the authorization code from the redirect URL
+            redirect_url = response.headers["location"]
+            parsed_url = urlparse(redirect_url)
+            query_params = parse_qs(parsed_url.query)
 
-        # Exchange the authorization code for tokens
-        response = await test_client.post(
-            "/token",
-            data={
-                "grant_type": "authorization_code",
-                "client_id": client_info["client_id"],
-                "client_secret": client_info["client_secret"],
-                "code": auth_code,
-                "code_verifier": pkce_challenge["code_verifier"],
-                "redirect_uri": "https://client.example.com/callback",
-            },
-        )
-        assert response.status_code == 200
+            assert "code" in query_params
+            auth_code = query_params["code"][0]
 
-        token_response = response.json()
-        assert "access_token" in token_response
-        authorization = f"Bearer {token_response['access_token']}"
-
-        # Test the authenticated endpoint with valid token
-        async with aconnect_sse(
-            test_client, "GET", "/sse", headers={"Authorization": authorization}
-        ) as event_source:
-            assert event_source.response.status_code == 200
-            events = event_source.aiter_sse()
-            sse = await events.__anext__()
-            assert sse.event == "endpoint"
-            assert sse.data.startswith("/messages/?session_id=")
-            messages_uri = sse.data
-
-            # verify that we can now post to the /messages endpoint, and get a response
-            # on the /sse endpoint
+            # Exchange the authorization code for tokens
             response = await test_client.post(
-                messages_uri,
-                headers={"Authorization": authorization},
-                content=JSONRPCRequest(
-                    jsonrpc="2.0",
-                    id="123",
-                    method="initialize",
-                    params={
-                        "protocolVersion": "2024-11-05",
-                        "capabilities": {
-                            "roots": {"listChanged": True},
-                            "sampling": {},
-                        },
-                        "clientInfo": {"name": "ExampleClient", "version": "1.0.0"},
-                    },
-                ).model_dump_json(),
-            )
-            assert response.status_code == 202
-            assert response.content == b"Accepted"
-
-            sse = await events.__anext__()
-            assert sse.event == "message"
-            sse_data = json.loads(sse.data)
-            assert sse_data["id"] == "123"
-            assert set(sse_data["result"]["capabilities"].keys()) == set(
-                ("experimental", "prompts", "resources", "tools")
+                "/token",
+                data={
+                    "grant_type": "authorization_code",
+                    "client_id": client_info["client_id"],
+                    "client_secret": client_info["client_secret"],
+                    "code": auth_code,
+                    "code_verifier": pkce_challenge["code_verifier"],
+                    "redirect_uri": "https://client.example.com/callback",
+                },
             )
+            assert response.status_code == 200
+
+            token_response = response.json()
+            assert "access_token" in token_response
+            authorization = f"Bearer {token_response['access_token']}"
+
+            # Test the authenticated endpoint with valid token
+            async with aconnect_sse(
+                test_client, "GET", "/sse", headers={"Authorization": authorization}
+            ) as event_source:
+                assert event_source.response.status_code == 200
+                events = event_source.aiter_sse()
+                sse = await events.__anext__()
+                assert sse.event == "endpoint"
+                assert sse.data.startswith("/messages/?session_id=")
+                messages_uri = sse.data
+
+                # verify that we can now post to the /messages endpoint, and get a response
+                # on the /sse endpoint
+                response = await test_client.post(
+                    messages_uri,
+                    headers={"Authorization": authorization},
+                    content=JSONRPCRequest(
+                        jsonrpc="2.0",
+                        id="123",
+                        method="initialize",
+                        params={
+                            "protocolVersion": "2024-11-05",
+                            "capabilities": {
+                                "roots": {"listChanged": True},
+                                "sampling": {},
+                            },
+                            "clientInfo": {"name": "ExampleClient", "version": "1.0.0"},
+                        },
+                    ).model_dump_json(),
+                )
+                assert response.status_code == 202
+                assert response.content == b"Accepted"
+
+                sse = await events.__anext__()
+                assert sse.event == "message"
+                sse_data = json.loads(sse.data)
+                assert sse_data["id"] == "123"
+                assert set(sse_data["result"]["capabilities"].keys()) == set(
+                    ("experimental", "prompts", "resources", "tools")
+                )
+                task_group.cancel_scope.cancel()
 
 
 class TestAuthorizeEndpointErrors: