diff --git a/tests/conftest.py b/tests/conftest.py index 94ddb0f9..a7a3ccfc 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -8,8 +8,9 @@ import subprocess import threading import time import uuid +from contextlib import asynccontextmanager from http.server import BaseHTTPRequestHandler, HTTPServer -from typing import Any, AsyncGenerator +from typing import Any, AsyncGenerator, AsyncIterator from urllib.parse import parse_qs, quote, urlparse import anyio @@ -118,6 +119,7 @@ async def wait_for_nextcloud( return False +@asynccontextmanager async def create_mcp_client_session( url: str, token: str | None = None, @@ -125,7 +127,7 @@ async def create_mcp_client_session( elicitation_callback: Any = None, sampling_callback: Any = None, headers: dict[str, str] | None = None, -) -> AsyncGenerator[ClientSession, Any]: +) -> AsyncIterator[ClientSession]: """ Factory function to create an MCP client session with proper lifecycle management. @@ -227,10 +229,10 @@ async def nc_mcp_client(anyio_backend) -> AsyncGenerator[ClientSession, Any]: Uses anyio pytest plugin for proper async fixture handling. """ - async for session in create_mcp_client_session( + async with create_mcp_client_session( url="http://localhost:8000/mcp", client_name="Basic MCP (HTTP)", - ): + ) as session: yield session @@ -246,11 +248,11 @@ async def nc_mcp_oauth_client( Uses headless browser automation suitable for CI/CD. Uses anyio pytest plugin for proper async fixture handling. """ - async for session in create_mcp_client_session( + async with create_mcp_client_session( url="http://localhost:8001/mcp", token=playwright_oauth_token, client_name="OAuth MCP (Playwright)", - ): + ) as session: yield session @@ -271,11 +273,11 @@ async def nc_mcp_basic_auth_client( credentials = base64.b64encode(b"admin:admin").decode("utf-8") auth_header = f"Basic {credentials}" - async for session in create_mcp_client_session( + async with create_mcp_client_session( url="http://localhost:8003/mcp", headers={"Authorization": auth_header}, client_name="BasicAuth MCP (Multi-User)", - ): + ) as session: yield session @@ -296,11 +298,11 @@ async def nc_mcp_oauth_jwt_client( Uses headless browser automation suitable for CI/CD. Uses anyio pytest plugin for proper async fixture handling. """ - async for session in create_mcp_client_session( + async with create_mcp_client_session( url="http://localhost:8001/mcp", token=playwright_oauth_token_jwt, client_name="OAuth JWT MCP (Playwright)", - ): + ) as session: yield session @@ -456,12 +458,12 @@ async def nc_mcp_oauth_client_with_elicitation( await page.close() # Create client session with elicitation callback - async for session in create_mcp_client_session( + async with create_mcp_client_session( url="http://localhost:8001/mcp", token=playwright_oauth_token, client_name="OAuth MCP with Elicitation", elicitation_callback=elicitation_callback, - ): + ) as session: # Attach elicitation metadata for test validation session.elicitation_triggered = elicitation_triggered yield session @@ -482,11 +484,11 @@ async def nc_mcp_oauth_client_read_only( Uses JWT tokens because they embed scope information in claims, enabling proper scope-based tool filtering. """ - async for session in create_mcp_client_session( + async with create_mcp_client_session( url="http://localhost:8001/mcp", token=playwright_oauth_token_read_only, client_name="OAuth JWT MCP Read-Only (Playwright)", - ): + ) as session: yield session @@ -505,11 +507,11 @@ async def nc_mcp_oauth_client_write_only( Uses JWT tokens because they embed scope information in claims, enabling proper scope-based tool filtering. """ - async for session in create_mcp_client_session( + async with create_mcp_client_session( url="http://localhost:8001/mcp", token=playwright_oauth_token_write_only, client_name="OAuth JWT MCP Write-Only (Playwright)", - ): + ) as session: yield session @@ -527,11 +529,11 @@ async def nc_mcp_oauth_client_full_access( Uses JWT tokens because they embed scope information in claims, enabling proper scope-based tool filtering. """ - async for session in create_mcp_client_session( + async with create_mcp_client_session( url="http://localhost:8001/mcp", token=playwright_oauth_token_full_access, client_name="OAuth JWT MCP Full Access (Playwright)", - ): + ) as session: yield session @@ -552,11 +554,11 @@ async def nc_mcp_oauth_client_no_custom_scopes( Uses JWT tokens because they embed scope information in claims, enabling proper scope-based tool filtering. """ - async for session in create_mcp_client_session( + async with create_mcp_client_session( url="http://localhost:8001/mcp", token=playwright_oauth_token_no_custom_scopes, client_name="OAuth JWT MCP No Custom Scopes (Playwright)", - ): + ) as session: yield session @@ -2726,11 +2728,11 @@ async def alice_mcp_client( alice_oauth_token: str, ) -> AsyncGenerator[ClientSession, Any]: """MCP client authenticated as alice (owner role).""" - async for session in create_mcp_client_session( + async with create_mcp_client_session( url="http://localhost:8001/mcp", token=alice_oauth_token, client_name="Alice MCP", - ): + ) as session: yield session @@ -2739,11 +2741,11 @@ async def bob_mcp_client( anyio_backend, bob_oauth_token: str ) -> AsyncGenerator[ClientSession, Any]: """MCP client authenticated as bob (viewer role).""" - async for session in create_mcp_client_session( + async with create_mcp_client_session( url="http://localhost:8001/mcp", token=bob_oauth_token, client_name="Bob MCP", - ): + ) as session: yield session @@ -2753,11 +2755,11 @@ async def charlie_mcp_client( charlie_oauth_token: str, ) -> AsyncGenerator[ClientSession, Any]: """MCP client authenticated as charlie (editor role, in 'editors' group).""" - async for session in create_mcp_client_session( + async with create_mcp_client_session( url="http://localhost:8001/mcp", token=charlie_oauth_token, client_name="Charlie MCP", - ): + ) as session: yield session @@ -2767,11 +2769,11 @@ async def diana_mcp_client( diana_oauth_token: str, ) -> AsyncGenerator[ClientSession, Any]: """MCP client authenticated as diana (no-access role).""" - async for session in create_mcp_client_session( + async with create_mcp_client_session( url="http://localhost:8001/mcp", token=diana_oauth_token, client_name="Diana MCP", - ): + ) as session: yield session diff --git a/tests/integration/sampling_support.py b/tests/integration/sampling_support.py index 4dae6999..20989a96 100644 --- a/tests/integration/sampling_support.py +++ b/tests/integration/sampling_support.py @@ -39,10 +39,10 @@ def create_sampling_callback(provider: Provider): if provider.supports_generation: callback = create_sampling_callback(provider) - async for session in create_mcp_client_session( + async with create_mcp_client_session( url="http://localhost:8000/mcp", sampling_callback=callback, - ): + ) as session: # Session now supports sampling pass ``` diff --git a/tests/integration/test_astrolabe_plotly_visualization.py b/tests/integration/test_astrolabe_plotly_visualization.py index 37b99443..e956e0c0 100644 --- a/tests/integration/test_astrolabe_plotly_visualization.py +++ b/tests/integration/test_astrolabe_plotly_visualization.py @@ -145,11 +145,11 @@ async def test_astrolabe_plotly_visualization_with_basic_auth( logger.info(f"Authorization result: {auth_result}") # Create MCP client session as alice - all MCP operations inside this block - async for alice_mcp_client in create_mcp_client_session( + async with create_mcp_client_session( url="http://localhost:8003/mcp", headers={"Authorization": auth_header}, client_name="Alice BasicAuth MCP", - ): + ) as alice_mcp_client: # Phase 3: Get initial indexed count initial_sync = await alice_mcp_client.call_tool( "nc_get_vector_sync_status", {} @@ -355,11 +355,11 @@ The visualization should show this document as a point in PCA-reduced space. # Cleanup note if not already cleaned (create new client for cleanup) if note_id: try: - async for cleanup_client in create_mcp_client_session( + async with create_mcp_client_session( url="http://localhost:8003/mcp", headers={"Authorization": auth_header}, client_name="Cleanup MCP", - ): + ) as cleanup_client: delete_response = await cleanup_client.call_tool( "nc_notes_delete_note", {"note_id": note_id} ) diff --git a/tests/integration/test_rag.py b/tests/integration/test_rag.py index 54996a2c..c9876fd2 100644 --- a/tests/integration/test_rag.py +++ b/tests/integration/test_rag.py @@ -227,11 +227,11 @@ async def nc_mcp_client_with_sampling( """ sampling_callback = create_sampling_callback(generation_provider) - async for session in create_mcp_client_session( + async with create_mcp_client_session( url="http://localhost:8000/mcp", client_name=f"Sampling MCP ({provider_name})", sampling_callback=sampling_callback, - ): + ) as session: yield session diff --git a/tests/server/login_flow/conftest.py b/tests/server/login_flow/conftest.py index 73c65c03..703ed79b 100644 --- a/tests/server/login_flow/conftest.py +++ b/tests/server/login_flow/conftest.py @@ -367,12 +367,12 @@ async def nc_mcp_login_flow_client( content={"acknowledged": True}, ) - async for session in create_mcp_client_session( + async with create_mcp_client_session( url=LOGIN_FLOW_MCP_URL, token=login_flow_oauth_token, client_name="Login Flow MCP", elicitation_callback=elicitation_callback, - ): + ) as session: # Step 1: Provision access via Login Flow v2 logger.info("Starting Login Flow v2 provisioning...") provision_result = await session.call_tool( @@ -499,11 +499,11 @@ async def nc_mcp_login_flow_client_read_only( anyio_backend, login_flow_read_only_token: str ) -> AsyncGenerator[ClientSession, Any]: """MCP client with read-only scopes on the login-flow server.""" - async for session in create_mcp_client_session( + async with create_mcp_client_session( url=LOGIN_FLOW_MCP_URL, token=login_flow_read_only_token, client_name="Login Flow MCP Read-Only", - ): + ) as session: yield session @@ -512,11 +512,11 @@ async def nc_mcp_login_flow_client_write_only( anyio_backend, login_flow_write_only_token: str ) -> AsyncGenerator[ClientSession, Any]: """MCP client with write-only scopes on the login-flow server.""" - async for session in create_mcp_client_session( + async with create_mcp_client_session( url=LOGIN_FLOW_MCP_URL, token=login_flow_write_only_token, client_name="Login Flow MCP Write-Only", - ): + ) as session: yield session @@ -525,11 +525,11 @@ async def nc_mcp_login_flow_client_full_access( anyio_backend, login_flow_full_access_token: str ) -> AsyncGenerator[ClientSession, Any]: """MCP client with full access scopes on the login-flow server.""" - async for session in create_mcp_client_session( + async with create_mcp_client_session( url=LOGIN_FLOW_MCP_URL, token=login_flow_full_access_token, client_name="Login Flow MCP Full Access", - ): + ) as session: yield session @@ -538,11 +538,11 @@ async def nc_mcp_login_flow_client_no_custom_scopes( anyio_backend, login_flow_no_custom_scopes_token: str ) -> AsyncGenerator[ClientSession, Any]: """MCP client with no custom scopes on the login-flow server.""" - async for session in create_mcp_client_session( + async with create_mcp_client_session( url=LOGIN_FLOW_MCP_URL, token=login_flow_no_custom_scopes_token, client_name="Login Flow MCP No Custom Scopes", - ): + ) as session: yield session @@ -724,12 +724,12 @@ async def _provision_login_flow_mcp_client( return ElicitResult(action="accept", content={"acknowledged": True}) - async for session in create_mcp_client_session( + async with create_mcp_client_session( url=LOGIN_FLOW_MCP_URL, token=token, client_name=f"Login Flow MCP ({username})", elicitation_callback=elicitation_callback, - ): + ) as session: # Provision access provision_result = await session.call_tool( "nc_auth_provision_access", {"scopes": None} diff --git a/third_party/astrolabe b/third_party/astrolabe index 42ad0193..2bf4a705 160000 --- a/third_party/astrolabe +++ b/third_party/astrolabe @@ -1 +1 @@ -Subproject commit 42ad0193700d934bd1fc5cf5a065421fae5b65e2 +Subproject commit 2bf4a70599c8000a9f6f67f16ee3950767250772