feat: Split read/write scopes into app:read/write scopes
This commit is contained in:
+129
-76
@@ -1,9 +1,9 @@
|
||||
import asyncio
|
||||
import logging
|
||||
import os
|
||||
import uuid
|
||||
from typing import Any, AsyncGenerator
|
||||
|
||||
import anyio
|
||||
import httpx
|
||||
import pytest
|
||||
from httpx import HTTPStatusError
|
||||
@@ -14,6 +14,48 @@ from nextcloud_mcp_server.client import NextcloudClient
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
# Default scopes for OAuth testing - all app-specific read/write scopes
|
||||
DEFAULT_FULL_SCOPES = (
|
||||
"openid profile email "
|
||||
"notes:read notes:write "
|
||||
"calendar:read calendar:write "
|
||||
"todo:read todo:write "
|
||||
"contacts:read contacts:write "
|
||||
"cookbook:read cookbook:write "
|
||||
"deck:read deck:write "
|
||||
"tables:read tables:write "
|
||||
"files:read files:write "
|
||||
"sharing:read sharing:write"
|
||||
)
|
||||
|
||||
# Read-only scopes (all read scopes across apps) - should match DEFAULT_FULL_SCOPES read portion
|
||||
DEFAULT_READ_SCOPES = (
|
||||
"openid profile email "
|
||||
"notes:read "
|
||||
"calendar:read "
|
||||
"todo:read "
|
||||
"contacts:read "
|
||||
"cookbook:read "
|
||||
"deck:read "
|
||||
"tables:read "
|
||||
"files:read "
|
||||
"sharing:read"
|
||||
)
|
||||
|
||||
# Write-only scopes (all write scopes across apps) - should match DEFAULT_FULL_SCOPES write portion
|
||||
DEFAULT_WRITE_SCOPES = (
|
||||
"openid profile email "
|
||||
"notes:write "
|
||||
"calendar:write "
|
||||
"todo:write "
|
||||
"contacts:write "
|
||||
"cookbook:write "
|
||||
"deck:write "
|
||||
"tables:write "
|
||||
"files:write "
|
||||
"sharing:write"
|
||||
)
|
||||
|
||||
|
||||
@pytest.fixture(scope="session")
|
||||
def anyio_backend():
|
||||
@@ -56,7 +98,7 @@ async def wait_for_nextcloud(
|
||||
logger.info(
|
||||
f"Nextcloud not ready yet, waiting {delay}s... (attempt {attempt}/{max_attempts})"
|
||||
)
|
||||
await asyncio.sleep(delay)
|
||||
await anyio.sleep(delay)
|
||||
|
||||
logger.error(
|
||||
f"Nextcloud server at {host} did not become ready after {max_attempts} attempts"
|
||||
@@ -191,18 +233,18 @@ async def nc_mcp_oauth_jwt_client(
|
||||
) -> AsyncGenerator[ClientSession, Any]:
|
||||
"""
|
||||
Fixture to create an MCP client session for JWT OAuth integration tests.
|
||||
Connects to the JWT OAuth-enabled MCP server on port 8002 with OAuth authentication.
|
||||
Connects to the OAuth-enabled MCP server on port 8001 with JWT token authentication.
|
||||
|
||||
This server uses JWT tokens (RFC 9068) instead of opaque tokens, enabling:
|
||||
- Token introspection via JWT signature verification
|
||||
Uses JWT tokens (RFC 9068) which provide:
|
||||
- Token validation via JWT signature verification (JWKS)
|
||||
- Scope information embedded in token claims
|
||||
- Offline token validation without userinfo endpoint
|
||||
- Faster validation without userinfo endpoint call
|
||||
|
||||
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(
|
||||
url="http://localhost:8002/mcp",
|
||||
url="http://localhost:8001/mcp",
|
||||
token=playwright_oauth_token_jwt,
|
||||
client_name="OAuth JWT MCP (Playwright)",
|
||||
):
|
||||
@@ -215,17 +257,17 @@ async def nc_mcp_oauth_client_read_only(
|
||||
playwright_oauth_token_read_only: str,
|
||||
) -> AsyncGenerator[ClientSession, Any]:
|
||||
"""
|
||||
Fixture to create an MCP client session with only nc:read scope.
|
||||
Connects to the JWT OAuth-enabled MCP server on port 8002.
|
||||
Fixture to create an MCP client session with only read scopes.
|
||||
Connects to the OAuth-enabled MCP server on port 8001.
|
||||
|
||||
This client should only see read tools and should get 403 errors
|
||||
when attempting to call write tools.
|
||||
|
||||
Uses JWT MCP server because JWT tokens embed scope information in claims,
|
||||
enabling proper scope-based filtering.
|
||||
Uses JWT tokens because they embed scope information in claims,
|
||||
enabling proper scope-based tool filtering.
|
||||
"""
|
||||
async for session in create_mcp_client_session(
|
||||
url="http://localhost:8002/mcp",
|
||||
url="http://localhost:8001/mcp",
|
||||
token=playwright_oauth_token_read_only,
|
||||
client_name="OAuth JWT MCP Read-Only (Playwright)",
|
||||
):
|
||||
@@ -238,17 +280,17 @@ async def nc_mcp_oauth_client_write_only(
|
||||
playwright_oauth_token_write_only: str,
|
||||
) -> AsyncGenerator[ClientSession, Any]:
|
||||
"""
|
||||
Fixture to create an MCP client session with only nc:write scope.
|
||||
Connects to the JWT OAuth-enabled MCP server on port 8002.
|
||||
Fixture to create an MCP client session with only write scopes.
|
||||
Connects to the OAuth-enabled MCP server on port 8001.
|
||||
|
||||
This client should only see write tools and should get 403 errors
|
||||
when attempting to call read tools.
|
||||
|
||||
Uses JWT MCP server because JWT tokens embed scope information in claims,
|
||||
enabling proper scope-based filtering.
|
||||
Uses JWT tokens because they embed scope information in claims,
|
||||
enabling proper scope-based tool filtering.
|
||||
"""
|
||||
async for session in create_mcp_client_session(
|
||||
url="http://localhost:8002/mcp",
|
||||
url="http://localhost:8001/mcp",
|
||||
token=playwright_oauth_token_write_only,
|
||||
client_name="OAuth JWT MCP Write-Only (Playwright)",
|
||||
):
|
||||
@@ -261,16 +303,16 @@ async def nc_mcp_oauth_client_full_access(
|
||||
playwright_oauth_token_full_access: str,
|
||||
) -> AsyncGenerator[ClientSession, Any]:
|
||||
"""
|
||||
Fixture to create an MCP client session with both nc:read and nc:write scopes.
|
||||
Connects to the JWT OAuth-enabled MCP server on port 8002.
|
||||
Fixture to create an MCP client session with both read and write scopes.
|
||||
Connects to the OAuth-enabled MCP server on port 8001.
|
||||
|
||||
This client should see all tools and be able to call all operations.
|
||||
|
||||
Uses JWT MCP server because JWT tokens embed scope information in claims,
|
||||
enabling proper scope-based filtering.
|
||||
Uses JWT tokens because they embed scope information in claims,
|
||||
enabling proper scope-based tool filtering.
|
||||
"""
|
||||
async for session in create_mcp_client_session(
|
||||
url="http://localhost:8002/mcp",
|
||||
url="http://localhost:8001/mcp",
|
||||
token=playwright_oauth_token_full_access,
|
||||
client_name="OAuth JWT MCP Full Access (Playwright)",
|
||||
):
|
||||
@@ -284,18 +326,18 @@ async def nc_mcp_oauth_client_no_custom_scopes(
|
||||
) -> AsyncGenerator[ClientSession, Any]:
|
||||
"""
|
||||
Fixture to create an MCP client session with NO custom scopes.
|
||||
Connects to the JWT OAuth-enabled MCP server on port 8002.
|
||||
Connects to the OAuth-enabled MCP server on port 8001.
|
||||
|
||||
This client has only OIDC default scopes (openid, profile, email) without
|
||||
application-specific scopes (nc:read, nc:write).
|
||||
application-specific scopes (notes:read, notes:write, etc.).
|
||||
|
||||
Expected behavior: Should see 0 tools (all tools require custom scopes).
|
||||
|
||||
Uses JWT MCP server because JWT tokens embed scope information in claims,
|
||||
enabling proper scope-based filtering.
|
||||
Uses JWT tokens because they embed scope information in claims,
|
||||
enabling proper scope-based tool filtering.
|
||||
"""
|
||||
async for session in create_mcp_client_session(
|
||||
url="http://localhost:8002/mcp",
|
||||
url="http://localhost:8001/mcp",
|
||||
token=playwright_oauth_token_no_custom_scopes,
|
||||
client_name="OAuth JWT MCP No Custom Scopes (Playwright)",
|
||||
):
|
||||
@@ -682,10 +724,9 @@ async def shared_calendar_2(
|
||||
try:
|
||||
# Wait for first calendar to fully initialize to avoid Nextcloud rate limiting
|
||||
# When creating multiple calendars rapidly, Nextcloud may not register them all
|
||||
import asyncio
|
||||
|
||||
logger.info("Waiting before creating second calendar to avoid rate limiting...")
|
||||
await asyncio.sleep(3) # Increased from 2 to 3 seconds
|
||||
await anyio.sleep(3) # Increased from 2 to 3 seconds
|
||||
|
||||
# Create a test calendar
|
||||
logger.info(f"Creating second shared test calendar: {calendar_name}")
|
||||
@@ -703,9 +744,8 @@ async def shared_calendar_2(
|
||||
|
||||
# Verify calendar was created by listing calendars
|
||||
# Add small delay to allow calendar to propagate in the system
|
||||
import asyncio
|
||||
|
||||
await asyncio.sleep(1.0) # Allow time for calendar to propagate
|
||||
await anyio.sleep(1.0) # Allow time for calendar to propagate
|
||||
|
||||
calendars = await nc_client.calendar.list_calendars()
|
||||
calendar_names = [cal["name"] for cal in calendars]
|
||||
@@ -714,7 +754,7 @@ async def shared_calendar_2(
|
||||
f"Calendar {calendar_name} not found immediately after creation. Available: {calendar_names}"
|
||||
)
|
||||
# Try one more time after a longer delay
|
||||
await asyncio.sleep(3) # Additional wait for calendar synchronization
|
||||
await anyio.sleep(3) # Additional wait for calendar synchronization
|
||||
calendars = await nc_client.calendar.list_calendars()
|
||||
calendar_names = [cal["name"] for cal in calendars]
|
||||
if calendar_name not in calendar_names:
|
||||
@@ -946,7 +986,7 @@ async def shared_oauth_client_credentials(anyio_backend, oauth_callback_server):
|
||||
client_id, client_secret = await _create_oauth_client_with_scopes(
|
||||
callback_url=callback_url,
|
||||
client_name="Pytest - Shared Test Client (Opaque)",
|
||||
allowed_scopes="openid profile email nc:read nc:write",
|
||||
allowed_scopes=DEFAULT_FULL_SCOPES,
|
||||
token_type="Bearer", # Opaque tokens for port 8001
|
||||
cache_file=".nextcloud_oauth_shared_test_client.json",
|
||||
)
|
||||
@@ -968,10 +1008,11 @@ async def shared_oauth_client_credentials(anyio_backend, oauth_callback_server):
|
||||
@pytest.fixture(scope="session")
|
||||
async def shared_jwt_oauth_client_credentials(anyio_backend, oauth_callback_server):
|
||||
"""
|
||||
Fixture to obtain shared JWT OAuth client credentials for JWT MCP server.
|
||||
Fixture to obtain shared JWT OAuth client credentials for testing JWT token behavior.
|
||||
|
||||
Creates a JWT OAuth client with full scopes (nc:read and nc:write) for use with
|
||||
the JWT MCP server (port 8002) that validates JWT tokens locally.
|
||||
Creates a JWT OAuth client with full scopes (all app read/write scopes). The client
|
||||
is configured with token_type="JWT" to request JWT-formatted access tokens from the
|
||||
OIDC server (instead of opaque tokens).
|
||||
|
||||
Returns:
|
||||
Tuple of (client_id, client_secret, callback_url, token_endpoint, authorization_endpoint)
|
||||
@@ -1001,12 +1042,12 @@ async def shared_jwt_oauth_client_credentials(anyio_backend, oauth_callback_serv
|
||||
"OIDC discovery missing required endpoints (token_endpoint or authorization_endpoint)"
|
||||
)
|
||||
|
||||
# Create JWT client with full scopes (nc:read and nc:write)
|
||||
# Create JWT client with full scopes (all app read/write scopes)
|
||||
# Cache to file to avoid creating new client on every test run
|
||||
client_id, client_secret = await _create_oauth_client_with_scopes(
|
||||
callback_url=callback_url,
|
||||
client_name="Pytest - Shared JWT Test Client",
|
||||
allowed_scopes="openid profile email nc:read nc:write",
|
||||
allowed_scopes=DEFAULT_FULL_SCOPES,
|
||||
token_type="JWT", # Explicitly set JWT token type
|
||||
cache_file=".nextcloud_oauth_shared_jwt_test_client.json",
|
||||
)
|
||||
@@ -1140,7 +1181,7 @@ async def _create_oauth_client_with_scopes(
|
||||
@pytest.fixture(scope="session")
|
||||
async def read_only_oauth_client_credentials(anyio_backend, oauth_callback_server):
|
||||
"""
|
||||
Fixture for OAuth client with only nc:read scope.
|
||||
Fixture for OAuth client with only read scopes.
|
||||
|
||||
Returns:
|
||||
Tuple of (client_id, client_secret, callback_url, token_endpoint, authorization_endpoint)
|
||||
@@ -1164,7 +1205,7 @@ async def read_only_oauth_client_credentials(anyio_backend, oauth_callback_serve
|
||||
client_id, client_secret = await _create_oauth_client_with_scopes(
|
||||
callback_url=callback_url,
|
||||
client_name="Test Client Read Only",
|
||||
allowed_scopes="openid profile email nc:read",
|
||||
allowed_scopes=DEFAULT_READ_SCOPES,
|
||||
token_type="JWT", # JWT tokens for scope validation
|
||||
)
|
||||
|
||||
@@ -1180,7 +1221,7 @@ async def read_only_oauth_client_credentials(anyio_backend, oauth_callback_serve
|
||||
@pytest.fixture(scope="session")
|
||||
async def write_only_oauth_client_credentials(anyio_backend, oauth_callback_server):
|
||||
"""
|
||||
Fixture for OAuth client with only nc:write scope.
|
||||
Fixture for OAuth client with only write scopes.
|
||||
|
||||
Returns:
|
||||
Tuple of (client_id, client_secret, callback_url, token_endpoint, authorization_endpoint)
|
||||
@@ -1204,7 +1245,7 @@ async def write_only_oauth_client_credentials(anyio_backend, oauth_callback_serv
|
||||
client_id, client_secret = await _create_oauth_client_with_scopes(
|
||||
callback_url=callback_url,
|
||||
client_name="Test Client Write Only",
|
||||
allowed_scopes="openid profile email nc:write",
|
||||
allowed_scopes=DEFAULT_WRITE_SCOPES,
|
||||
token_type="JWT", # JWT tokens for scope validation
|
||||
)
|
||||
|
||||
@@ -1220,7 +1261,7 @@ async def write_only_oauth_client_credentials(anyio_backend, oauth_callback_serv
|
||||
@pytest.fixture(scope="session")
|
||||
async def full_access_oauth_client_credentials(anyio_backend, oauth_callback_server):
|
||||
"""
|
||||
Fixture for OAuth client with both nc:read and nc:write scopes.
|
||||
Fixture for OAuth client with both read and write scopes.
|
||||
|
||||
Returns:
|
||||
Tuple of (client_id, client_secret, callback_url, token_endpoint, authorization_endpoint)
|
||||
@@ -1244,7 +1285,7 @@ async def full_access_oauth_client_credentials(anyio_backend, oauth_callback_ser
|
||||
client_id, client_secret = await _create_oauth_client_with_scopes(
|
||||
callback_url=callback_url,
|
||||
client_name="Test Client Full Access",
|
||||
allowed_scopes="openid profile email nc:read nc:write",
|
||||
allowed_scopes=DEFAULT_FULL_SCOPES,
|
||||
token_type="JWT", # JWT tokens for scope validation
|
||||
)
|
||||
|
||||
@@ -1265,7 +1306,7 @@ async def no_custom_scopes_oauth_client_credentials(
|
||||
Fixture for OAuth client with NO custom scopes (only OIDC defaults).
|
||||
|
||||
Tests the security behavior when a user grants only the default OIDC scopes
|
||||
(openid, profile, email) but declines custom application scopes (nc:read, nc:write).
|
||||
(openid, profile, email) but declines custom application scopes (notes:read, notes:write, etc.).
|
||||
|
||||
Returns:
|
||||
Tuple of (client_id, client_secret, callback_url, token_endpoint, authorization_endpoint)
|
||||
@@ -1289,7 +1330,7 @@ async def no_custom_scopes_oauth_client_credentials(
|
||||
client_id, client_secret = await _create_oauth_client_with_scopes(
|
||||
callback_url=callback_url,
|
||||
client_name="Test Client No Custom Scopes",
|
||||
allowed_scopes="openid profile email", # No nc:read or nc:write
|
||||
allowed_scopes="openid profile email", # No app-specific scopes (no app access)
|
||||
token_type="JWT", # JWT tokens for scope validation
|
||||
)
|
||||
|
||||
@@ -1363,7 +1404,7 @@ async def playwright_oauth_token(
|
||||
f"client_id={client_id}&"
|
||||
f"redirect_uri={quote(callback_url, safe='')}&"
|
||||
f"state={state}&"
|
||||
f"scope=openid%20profile%20email%20nc:read%20nc:write"
|
||||
f"scope=openid%20profile%20email%20notes:read%20notes:write%20calendar:read%20calendar:write%20contacts:read%20contacts:write%20cookbook:read%20cookbook:write%20deck:read%20deck:write%20tables:read%20tables:write%20files:read%20files:write%20sharing:read%20sharing:write"
|
||||
)
|
||||
|
||||
# Async browser automation using pytest-playwright's browser fixture
|
||||
@@ -1420,7 +1461,7 @@ async def playwright_oauth_token(
|
||||
raise TimeoutError(
|
||||
f"Timeout waiting for OAuth callback (state={state[:16]}...)"
|
||||
)
|
||||
await asyncio.sleep(0.5)
|
||||
await anyio.sleep(0.5)
|
||||
|
||||
auth_code = auth_states[state]
|
||||
logger.info(f"Successfully received authorization code: {auth_code[:20]}...")
|
||||
@@ -1460,7 +1501,7 @@ async def playwright_oauth_token_jwt(
|
||||
"""
|
||||
Fixture to obtain a JWT OAuth access token for the JWT MCP server.
|
||||
|
||||
Uses a JWT OAuth client with full scopes (nc:read and nc:write) to ensure
|
||||
Uses a JWT OAuth client with full scopes (all app read/write scopes) to ensure
|
||||
the access token includes proper scope claims that the JWT MCP server can validate.
|
||||
|
||||
Returns:
|
||||
@@ -1470,7 +1511,7 @@ async def playwright_oauth_token_jwt(
|
||||
browser,
|
||||
shared_jwt_oauth_client_credentials,
|
||||
oauth_callback_server,
|
||||
scopes="openid profile email nc:read nc:write",
|
||||
scopes=DEFAULT_FULL_SCOPES,
|
||||
)
|
||||
|
||||
|
||||
@@ -1591,7 +1632,7 @@ async def _get_oauth_token_with_scopes(
|
||||
browser: Playwright browser instance
|
||||
shared_oauth_client_credentials: Tuple of OAuth client credentials
|
||||
oauth_callback_server: OAuth callback server fixture
|
||||
scopes: Space-separated list of scopes (e.g., "openid profile email nc:read")
|
||||
scopes: Space-separated list of scopes (e.g., "openid profile email notes:read")
|
||||
|
||||
Returns:
|
||||
OAuth access token string with requested scopes
|
||||
@@ -1688,7 +1729,7 @@ async def _get_oauth_token_with_scopes(
|
||||
auth_code = auth_states[state]
|
||||
logger.info("Auth code received from callback server")
|
||||
break
|
||||
await asyncio.sleep(0.1)
|
||||
await anyio.sleep(0.1)
|
||||
else:
|
||||
raise TimeoutError(
|
||||
f"Auth code not received within {timeout}s. State: {state[:16]}..."
|
||||
@@ -1727,18 +1768,18 @@ async def playwright_oauth_token_read_only(
|
||||
anyio_backend, browser, read_only_oauth_client_credentials, oauth_callback_server
|
||||
) -> str:
|
||||
"""
|
||||
Fixture to obtain an OAuth access token with only nc:read scope.
|
||||
Fixture to obtain an OAuth access token with only read scopes.
|
||||
|
||||
This token will only be able to perform read operations and should
|
||||
have write tools filtered out from the tool list.
|
||||
|
||||
Uses a dedicated OAuth client with allowed_scopes="openid profile email nc:read"
|
||||
Uses a dedicated OAuth client with allowed_scopes=DEFAULT_READ_SCOPES
|
||||
"""
|
||||
return await _get_oauth_token_with_scopes(
|
||||
browser,
|
||||
read_only_oauth_client_credentials,
|
||||
oauth_callback_server,
|
||||
scopes="openid profile email nc:read",
|
||||
scopes=DEFAULT_READ_SCOPES,
|
||||
)
|
||||
|
||||
|
||||
@@ -1747,18 +1788,18 @@ async def playwright_oauth_token_write_only(
|
||||
anyio_backend, browser, write_only_oauth_client_credentials, oauth_callback_server
|
||||
) -> str:
|
||||
"""
|
||||
Fixture to obtain an OAuth access token with only nc:write scope.
|
||||
Fixture to obtain an OAuth access token with only write scopes.
|
||||
|
||||
This token will only be able to perform write operations and should
|
||||
have read tools filtered out from the tool list.
|
||||
|
||||
Uses a dedicated OAuth client with allowed_scopes="openid profile email nc:write"
|
||||
Uses a dedicated OAuth client with allowed_scopes=DEFAULT_WRITE_SCOPES
|
||||
"""
|
||||
return await _get_oauth_token_with_scopes(
|
||||
browser,
|
||||
write_only_oauth_client_credentials,
|
||||
oauth_callback_server,
|
||||
scopes="openid profile email nc:write",
|
||||
scopes=DEFAULT_WRITE_SCOPES,
|
||||
)
|
||||
|
||||
|
||||
@@ -1767,17 +1808,17 @@ async def playwright_oauth_token_full_access(
|
||||
anyio_backend, browser, full_access_oauth_client_credentials, oauth_callback_server
|
||||
) -> str:
|
||||
"""
|
||||
Fixture to obtain an OAuth access token with both nc:read and nc:write scopes.
|
||||
Fixture to obtain an OAuth access token with both read and write scopes.
|
||||
|
||||
This token will be able to perform all operations.
|
||||
|
||||
Uses a dedicated JWT OAuth client with allowed_scopes="openid profile email nc:read nc:write"
|
||||
Uses a dedicated JWT OAuth client with allowed_scopes=DEFAULT_FULL_SCOPES
|
||||
"""
|
||||
return await _get_oauth_token_with_scopes(
|
||||
browser,
|
||||
full_access_oauth_client_credentials,
|
||||
oauth_callback_server,
|
||||
scopes="openid profile email nc:read nc:write",
|
||||
scopes=DEFAULT_FULL_SCOPES,
|
||||
)
|
||||
|
||||
|
||||
@@ -1795,7 +1836,7 @@ async def playwright_oauth_token_no_custom_scopes(
|
||||
(openid, profile, email) but declines application-specific scopes.
|
||||
|
||||
Expected: JWT token will contain only default scopes, and all MCP tools
|
||||
should be filtered out since they all require nc:read or nc:write.
|
||||
should be filtered out since they all require app-specific scopes.
|
||||
|
||||
Uses a dedicated JWT OAuth client with allowed_scopes="openid profile email"
|
||||
"""
|
||||
@@ -1967,7 +2008,7 @@ async def _get_oauth_token_for_user(
|
||||
f"client_id={client_id}&"
|
||||
f"redirect_uri={quote(callback_url, safe='')}&"
|
||||
f"state={state}&"
|
||||
f"scope=openid%20profile%20email%20nc:read%20nc:write"
|
||||
f"scope=openid%20profile%20email%20notes:read%20notes:write%20calendar:read%20calendar:write%20contacts:read%20contacts:write%20cookbook:read%20cookbook:write%20deck:read%20deck:write%20tables:read%20tables:write%20files:read%20files:write%20sharing:read%20sharing:write"
|
||||
)
|
||||
|
||||
logger.info(f"Performing browser OAuth flow for {username}...")
|
||||
@@ -2013,7 +2054,7 @@ async def _get_oauth_token_for_user(
|
||||
raise TimeoutError(
|
||||
f"Timeout waiting for OAuth callback for {username} (state={state[:16]}...)"
|
||||
)
|
||||
await asyncio.sleep(0.5)
|
||||
await anyio.sleep(0.5)
|
||||
|
||||
auth_code = auth_states[state]
|
||||
logger.info(f"Got auth code for {username}: {auth_code[:20]}...")
|
||||
@@ -2064,7 +2105,6 @@ async def all_oauth_tokens(
|
||||
Now uses the real callback server with state parameters for reliable
|
||||
concurrent token acquisition without race conditions.
|
||||
"""
|
||||
import asyncio
|
||||
import time
|
||||
|
||||
# Get auth_states dict from callback server
|
||||
@@ -2077,7 +2117,7 @@ async def all_oauth_tokens(
|
||||
async def get_token_with_delay(username: str, config: dict, delay: float):
|
||||
"""Get token for a user after a small delay to stagger requests."""
|
||||
if delay > 0:
|
||||
await asyncio.sleep(delay)
|
||||
await anyio.sleep(delay)
|
||||
return await _get_oauth_token_for_user(
|
||||
browser,
|
||||
shared_oauth_client_credentials,
|
||||
@@ -2087,17 +2127,30 @@ async def all_oauth_tokens(
|
||||
)
|
||||
|
||||
# Create tasks for all users with staggered starts (0.5s apart)
|
||||
tasks = {
|
||||
username: get_token_with_delay(username, config, idx * 0.5)
|
||||
for idx, (username, config) in enumerate(test_users_setup.items())
|
||||
}
|
||||
|
||||
# Run all token fetches concurrently
|
||||
results = await asyncio.gather(*tasks.values(), return_exceptions=True)
|
||||
|
||||
# Build result dict, handling any errors
|
||||
user_list = list(test_users_setup.items())
|
||||
tokens = {}
|
||||
for username, result in zip(tasks.keys(), results):
|
||||
|
||||
# Run all token fetches concurrently using anyio task groups
|
||||
async with anyio.create_task_group() as tg:
|
||||
# Create a dict to store results as they complete
|
||||
results = {}
|
||||
|
||||
def create_task_wrapper(username: str, config: dict, idx: int):
|
||||
async def task():
|
||||
try:
|
||||
token = await get_token_with_delay(username, config, idx * 0.5)
|
||||
results[username] = token
|
||||
except Exception as e:
|
||||
results[username] = e
|
||||
|
||||
return task
|
||||
|
||||
for idx, (username, config) in enumerate(user_list):
|
||||
tg.start_soon(create_task_wrapper(username, config, idx))
|
||||
|
||||
# Build token dict, handling any errors
|
||||
for username in results:
|
||||
result = results[username]
|
||||
if isinstance(result, Exception):
|
||||
logger.error(f"Failed to get OAuth token for {username}: {result}")
|
||||
raise result
|
||||
|
||||
Reference in New Issue
Block a user