From 240ceb38087c330dd19bd6f254ecf2ca153618b9 Mon Sep 17 00:00:00 2001
From: Chris Coutinho <chris@coutinho.io>
Date: Sat, 18 Oct 2025 21:03:24 +0200
Subject: [PATCH] test: Migrate load test framework to anyio as well

---
 tests/load/benchmark.py              |  3 ++-
 tests/load/cleanup_loadtest_users.py |  4 ++--
 tests/load/oauth_benchmark.py        |  3 ++-
 tests/load/oauth_pool.py             | 25 ++-----------------------
 4 files changed, 8 insertions(+), 27 deletions(-)

diff --git a/tests/load/benchmark.py b/tests/load/benchmark.py
index 020bebb..54adffb 100644
--- a/tests/load/benchmark.py
+++ b/tests/load/benchmark.py
@@ -18,6 +18,7 @@ from collections import Counter
 from contextlib import asynccontextmanager
 from typing import Any
 
+import anyio
 import click
 from mcp import ClientSession
 from mcp.client.streamable_http import streamablehttp_client
@@ -494,7 +495,7 @@ def main(
             print(f"Results exported to: {output}")
 
     try:
-        asyncio.run(run())
+        anyio.run(run)
     except KeyboardInterrupt:
         print("\nBenchmark interrupted by user")
         sys.exit(130)
diff --git a/tests/load/cleanup_loadtest_users.py b/tests/load/cleanup_loadtest_users.py
index b233faf..1492b23 100644
--- a/tests/load/cleanup_loadtest_users.py
+++ b/tests/load/cleanup_loadtest_users.py
@@ -11,9 +11,9 @@ Usage:
     uv run python -m tests.load.cleanup_loadtest_users --dry-run
 """
 
-import asyncio
 import sys
 
+import anyio
 import click
 
 from nextcloud_mcp_server.client import NextcloudClient
@@ -110,7 +110,7 @@ def main(prefix: str, dry_run: bool):
         # Delete users with custom prefix
         uv run python -m tests.load.cleanup_loadtest_users --prefix mytest
     """
-    asyncio.run(cleanup_users(prefix=prefix, dry_run=dry_run))
+    anyio.run(cleanup_users, prefix, dry_run)
 
 
 if __name__ == "__main__":
diff --git a/tests/load/oauth_benchmark.py b/tests/load/oauth_benchmark.py
index 56505ad..4cf3296 100644
--- a/tests/load/oauth_benchmark.py
+++ b/tests/load/oauth_benchmark.py
@@ -23,6 +23,7 @@ from http.server import BaseHTTPRequestHandler, HTTPServer
 from typing import Any
 from urllib.parse import parse_qs, urlparse
 
+import anyio
 import click
 import httpx
 from playwright.async_api import async_playwright
@@ -729,7 +730,7 @@ def main(
             print(f"Results exported to: {output}")
 
     try:
-        asyncio.run(run())
+        anyio.run(run)
     except KeyboardInterrupt:
         print("\nBenchmark interrupted by user")
         sys.exit(130)
diff --git a/tests/load/oauth_pool.py b/tests/load/oauth_pool.py
index 3d1eaea..9ed4fea 100644
--- a/tests/load/oauth_pool.py
+++ b/tests/load/oauth_pool.py
@@ -180,13 +180,8 @@ class OAuthUserPool:
             # Clean up streamable context if session creation failed
             try:
                 await streamable_context.__aexit__(None, None, None)
-            except RuntimeError as cleanup_error:
-                if "cancel scope" in str(cleanup_error):
-                    logger.debug(
-                        f"Ignoring cancel scope teardown issue: {cleanup_error}"
-                    )
-                else:
-                    raise
+            except Exception as cleanup_error:
+                logger.debug(f"Error during cleanup: {cleanup_error}")
             raise e
 
     async def close_user_session(self, username: str):
@@ -200,13 +195,6 @@ class OAuthUserPool:
         if profile.session:
             try:
                 await profile.session.__aexit__(None, None, None)
-            except RuntimeError as e:
-                if "cancel scope" in str(e):
-                    logger.debug(
-                        f"Ignoring cancel scope teardown issue for {username}: {e}"
-                    )
-                else:
-                    logger.debug(f"Error closing session for {username}: {e}")
             except Exception as e:
                 logger.debug(f"Error closing session for {username}: {e}")
             profile.session = None
@@ -215,15 +203,6 @@ class OAuthUserPool:
         if profile.streamable_context:
             try:
                 await profile.streamable_context.__aexit__(None, None, None)
-            except RuntimeError as e:
-                if "cancel scope" in str(e):
-                    logger.debug(
-                        f"Ignoring cancel scope teardown issue for {username}: {e}"
-                    )
-                else:
-                    logger.debug(
-                        f"Error closing streamable context for {username}: {e}"
-                    )
             except Exception as e:
                 logger.debug(f"Error closing streamable context for {username}: {e}")
             profile.streamable_context = None