Merge pull request #2 from open-webui/refac-entrypoint

refac: move env loading and user auth to benchmarking entrypoint used for all scripts

.env.example

@@ -10,9 +10,10 @@ OLLAMA_BASE_URL=http://host.docker.internal:11434
 # Enable Channels feature
 ENABLE_CHANNELS=true
 
-# Admin credentials for creating test channels and users
+# Admin credentials for Open WebUI (REQUIRED)
 ADMIN_USER_EMAIL=admin@example.com
 ADMIN_USER_PASSWORD=adminpassword123
+ADMIN_USER_NAME=Admin User
 
 # Benchmark settings (optional)
 MAX_CONCURRENT_USERS=50

benchmark/auth/__init__.py

@@ -0,0 +1,20 @@
+"""
+Authentication module for Open WebUI Benchmark.
+
+Provides consistent authentication handling across all benchmark scripts.
+"""
+
+from benchmark.auth.authenticator import Authenticator, AuthResult
+from benchmark.auth.entrypoint import (
+    ensure_admin_authenticated,
+    AuthenticationError,
+    ServiceNotReadyError,
+)
+
+__all__ = [
+    "Authenticator",
+    "AuthResult",
+    "ensure_admin_authenticated",
+    "AuthenticationError",
+    "ServiceNotReadyError",
+]

benchmark/auth/authenticator.py

@@ -0,0 +1,216 @@
+"""
+Authenticator class for Open WebUI.
+
+Encapsulates the signin-with-fallback-to-signup pattern for admin users.
+"""
+
+import os
+from dataclasses import dataclass
+from typing import Optional
+
+from dotenv import load_dotenv
+
+from benchmark.clients.http_client import OpenWebUIClient, User
+
+
+@dataclass
+class AuthResult:
+    """Result of an authentication attempt."""
+    success: bool
+    user: Optional[User] = None
+    client: Optional[OpenWebUIClient] = None
+    error: Optional[str] = None
+    is_new_signup: bool = False
+
+
+@dataclass
+class AdminCredentials:
+    """Admin user credentials loaded from environment."""
+    email: str
+    password: str
+    name: str = "Admin User"
+    
+    @classmethod
+    def from_env(cls) -> Optional["AdminCredentials"]:
+        """
+        Load admin credentials from environment variables.
+        
+        Looks for:
+        - ADMIN_USER_EMAIL
+        - ADMIN_USER_PASSWORD
+        - ADMIN_USER_NAME (optional, defaults to "Admin User")
+        
+        Returns:
+            AdminCredentials if both email and password are set, None otherwise.
+        """
+        load_dotenv()
+        
+        email = os.environ.get("ADMIN_USER_EMAIL")
+        password = os.environ.get("ADMIN_USER_PASSWORD")
+        name = os.environ.get("ADMIN_USER_NAME", "Admin User")
+        
+        if email and password:
+            return cls(email=email, password=password, name=name)
+        return None
+
+
+class Authenticator:
+    """
+    Handles authentication to Open WebUI instances.
+    
+    Provides a consistent signin-with-fallback-to-signup pattern
+    for admin user authentication.
+    """
+    
+    def __init__(
+        self,
+        base_url: str,
+        timeout: float = 30.0,
+    ):
+        """
+        Initialize the authenticator.
+        
+        Args:
+            base_url: Base URL of the Open WebUI instance
+            timeout: Request timeout in seconds
+        """
+        self.base_url = base_url.rstrip("/")
+        self.timeout = timeout
+        self._client: Optional[OpenWebUIClient] = None
+        self._authenticated_user: Optional[User] = None
+    
+    @property
+    def client(self) -> Optional[OpenWebUIClient]:
+        """Get the authenticated client."""
+        return self._client
+    
+    @property
+    def user(self) -> Optional[User]:
+        """Get the authenticated user."""
+        return self._authenticated_user
+    
+    @property
+    def token(self) -> Optional[str]:
+        """Get the authentication token."""
+        if self._authenticated_user:
+            return self._authenticated_user.token
+        return None
+    
+    @property
+    def is_authenticated(self) -> bool:
+        """Check if currently authenticated."""
+        return self._authenticated_user is not None
+    
+    async def _create_client(self) -> OpenWebUIClient:
+        """Create and connect an HTTP client."""
+        client = OpenWebUIClient(self.base_url, self.timeout)
+        await client.connect()
+        return client
+    
+    async def wait_for_service(
+        self,
+        max_retries: int = 30,
+        retry_delay: float = 1.0,
+    ) -> bool:
+        """
+        Wait for the Open WebUI service to be ready.
+        
+        Args:
+            max_retries: Maximum number of health check attempts
+            retry_delay: Delay between retries in seconds
+            
+        Returns:
+            True if service is ready, False otherwise
+        """
+        if self._client is None:
+            self._client = await self._create_client()
+        
+        # Convert to timeout/interval format expected by wait_for_ready
+        timeout = max_retries * retry_delay
+        return await self._client.wait_for_ready(
+            timeout=timeout,
+            interval=retry_delay,
+        )
+    
+    async def authenticate_admin(
+        self,
+        credentials: Optional[AdminCredentials] = None,
+    ) -> AuthResult:
+        """
+        Authenticate as admin user using signin with fallback to signup.
+        
+        If the admin user doesn't exist (first run), creates the account.
+        If the admin user exists, signs in with provided credentials.
+        
+        Args:
+            credentials: Admin credentials. If None, loads from environment.
+            
+        Returns:
+            AuthResult with authentication status and user details.
+        """
+        # Load credentials from environment if not provided
+        if credentials is None:
+            credentials = AdminCredentials.from_env()
+        
+        if credentials is None:
+            return AuthResult(
+                success=False,
+                error=(
+                    "Admin credentials not configured. "
+                    "Set ADMIN_USER_EMAIL and ADMIN_USER_PASSWORD environment variables."
+                ),
+            )
+        
+        # Create client if needed
+        if self._client is None:
+            self._client = await self._create_client()
+        
+        # Try signin first (existing user)
+        is_new_signup = False
+        try:
+            user = await self._client.signin(credentials.email, credentials.password)
+            self._authenticated_user = user
+            return AuthResult(
+                success=True,
+                user=user,
+                client=self._client,
+                is_new_signup=False,
+            )
+        except Exception as signin_error:
+            # Signin failed - try signup (new instance, first admin)
+            try:
+                user = await self._client.signup(
+                    credentials.email,
+                    credentials.password,
+                    credentials.name,
+                )
+                self._authenticated_user = user
+                return AuthResult(
+                    success=True,
+                    user=user,
+                    client=self._client,
+                    is_new_signup=True,
+                )
+            except Exception as signup_error:
+                return AuthResult(
+                    success=False,
+                    error=(
+                        f"Failed to authenticate admin user ({credentials.email}). "
+                        f"Signin error: {signin_error}. Signup error: {signup_error}"
+                    ),
+                )
+    
+    async def close(self) -> None:
+        """Close the client connection."""
+        if self._client:
+            await self._client.close()
+            self._client = None
+        self._authenticated_user = None
+    
+    async def __aenter__(self) -> "Authenticator":
+        """Async context manager entry."""
+        return self
+    
+    async def __aexit__(self, exc_type, exc_val, exc_tb) -> None:
+        """Async context manager exit."""
+        await self.close()

benchmark/auth/entrypoint.py

@@ -0,0 +1,185 @@
+"""
+Entrypoint functions for benchmark authentication.
+
+Provides high-level functions to ensure authentication is complete
+before any benchmark operations run.
+"""
+
+import os
+from typing import Optional
+
+from dotenv import load_dotenv
+
+from benchmark.auth.authenticator import Authenticator, AuthResult, AdminCredentials
+from benchmark.clients.http_client import OpenWebUIClient
+
+
+class AuthenticationError(Exception):
+    """Raised when authentication fails."""
+    pass
+
+
+class ServiceNotReadyError(Exception):
+    """Raised when the Open WebUI service is not reachable."""
+    pass
+
+
+async def ensure_admin_authenticated(
+    base_url: Optional[str] = None,
+    email: Optional[str] = None,
+    password: Optional[str] = None,
+    name: Optional[str] = None,
+    timeout: float = 30.0,
+    wait_for_service: bool = True,
+    service_wait_retries: int = 30,
+    service_retry_delay: float = 1.0,
+) -> tuple[OpenWebUIClient, AuthResult]:
+    """
+    Ensure admin user is authenticated before running benchmarks.
+    
+    This is the main entrypoint function that should be called before
+    any benchmark operations. It handles:
+    
+    1. Loading credentials from environment variables (if not provided)
+    2. Waiting for the Open WebUI service to be ready
+    3. Attempting signin with fallback to signup for first-time setup
+    
+    Args:
+        base_url: Target Open WebUI URL. If None, loads from OPEN_WEBUI_URL
+                  or BENCHMARK_TARGET_URL environment variable.
+        email: Admin email. If None, loads from ADMIN_USER_EMAIL env var.
+        password: Admin password. If None, loads from ADMIN_USER_PASSWORD env var.
+        name: Admin display name. If None, loads from ADMIN_USER_NAME or defaults.
+        timeout: HTTP request timeout in seconds.
+        wait_for_service: Whether to wait for service health check.
+        service_wait_retries: Max retries for service health check.
+        service_retry_delay: Delay between health check retries.
+        
+    Returns:
+        Tuple of (authenticated OpenWebUIClient, AuthResult)
+        
+    Raises:
+        ServiceNotReadyError: If service is not reachable after retries.
+        AuthenticationError: If authentication fails.
+        
+    Example:
+        ```python
+        from benchmark.auth import ensure_admin_authenticated
+        
+        async def run_my_benchmark():
+            client, auth_result = await ensure_admin_authenticated()
+            
+            if auth_result.is_new_signup:
+                print("Created new admin account")
+            
+            # Use authenticated client for benchmark operations
+            channels = await client.list_channels()
+            ...
+        ```
+    """
+    # Load environment variables
+    load_dotenv()
+    
+    # Resolve base URL
+    if base_url is None:
+        base_url = os.environ.get(
+            "OPEN_WEBUI_URL",
+            os.environ.get("BENCHMARK_TARGET_URL", "http://localhost:3000")
+        )
+    
+    # Build credentials
+    credentials = None
+    if email and password:
+        credentials = AdminCredentials(
+            email=email,
+            password=password,
+            name=name or "Admin User",
+        )
+    # Otherwise, Authenticator will load from env
+    
+    # Create authenticator
+    authenticator = Authenticator(base_url, timeout)
+    
+    try:
+        # Wait for service to be ready
+        if wait_for_service:
+            is_ready = await authenticator.wait_for_service(
+                max_retries=service_wait_retries,
+                retry_delay=service_retry_delay,
+            )
+            if not is_ready:
+                await authenticator.close()
+                raise ServiceNotReadyError(
+                    f"Open WebUI service at {base_url} is not ready after "
+                    f"{service_wait_retries} attempts."
+                )
+        
+        # Authenticate admin
+        result = await authenticator.authenticate_admin(credentials)
+        
+        if not result.success:
+            await authenticator.close()
+            raise AuthenticationError(result.error)
+        
+        # Return the client (don't close - caller owns it now)
+        return result.client, result
+        
+    except (ServiceNotReadyError, AuthenticationError):
+        # Re-raise our custom exceptions
+        raise
+    except Exception as e:
+        await authenticator.close()
+        raise AuthenticationError(f"Unexpected error during authentication: {e}")
+
+
+async def check_auth_status(
+    base_url: Optional[str] = None,
+    timeout: float = 10.0,
+) -> dict:
+    """
+    Check authentication configuration and service status.
+    
+    Useful for validating environment setup before running benchmarks.
+    Does not perform actual authentication.
+    
+    Args:
+        base_url: Target Open WebUI URL. If None, loads from environment.
+        timeout: HTTP request timeout in seconds.
+        
+    Returns:
+        Dictionary with status information:
+        - service_url: The resolved service URL
+        - service_reachable: Whether health endpoint responds
+        - credentials_configured: Whether admin credentials are in environment
+        - admin_email: The configured admin email (if set)
+    """
+    load_dotenv()
+    
+    # Resolve base URL
+    if base_url is None:
+        base_url = os.environ.get(
+            "OPEN_WEBUI_URL",
+            os.environ.get("BENCHMARK_TARGET_URL", "http://localhost:3000")
+        )
+    
+    # Check credentials
+    credentials = AdminCredentials.from_env()
+    
+    # Check service health
+    service_reachable = False
+    try:
+        authenticator = Authenticator(base_url, timeout)
+        service_reachable = await authenticator.wait_for_service(
+            max_retries=3,
+            retry_delay=1.0,
+        )
+        await authenticator.close()
+    except Exception:
+        pass
+    
+    return {
+        "service_url": base_url,
+        "service_reachable": service_reachable,
+        "credentials_configured": credentials is not None,
+        "admin_email": credentials.email if credentials else None,
+    }

benchmark/cli.py

@@ -16,6 +16,12 @@ from rich.panel import Panel
 from benchmark.core.config import load_config, ConfigLoader
 from benchmark.core.runner import BenchmarkRunner
 from benchmark.scenarios.channels import ChannelConcurrencyBenchmark, ChannelWebSocketBenchmark
+from benchmark.auth import (
+    ensure_admin_authenticated,
+    AuthenticationError,
+    ServiceNotReadyError,
+)
+from benchmark.auth.entrypoint import check_auth_status
 
 
 console = Console()
@@ -47,6 +53,76 @@ def list_profiles():
         console.print()
 
 
+async def auth_check(target_url: Optional[str] = None):
+    """Check authentication configuration and service status."""
+    console.print("\n[bold]Authentication Status Check[/bold]\n")
+    
+    with console.status("Checking configuration and service..."):
+        status = await check_auth_status(base_url=target_url)
+    
+    # Display results
+    console.print(f"  Service URL: [cyan]{status['service_url']}[/cyan]")
+    
+    if status['service_reachable']:
+        console.print(f"  Service Status: [green]✓ Reachable[/green]")
+    else:
+        console.print(f"  Service Status: [red]✗ Not reachable[/red]")
+    
+    if status['credentials_configured']:
+        console.print(f"  Credentials: [green]✓ Configured[/green]")
+        console.print(f"  Admin Email: [cyan]{status['admin_email']}[/cyan]")
+    else:
+        console.print(f"  Credentials: [red]✗ Not configured[/red]")
+        console.print("    Set ADMIN_USER_EMAIL and ADMIN_USER_PASSWORD environment variables")
+    
+    console.print()
+    
+    # Summary
+    if status['service_reachable'] and status['credentials_configured']:
+        console.print("[green]Ready to run benchmarks![/green]")
+        return True
+    else:
+        console.print("[yellow]Configuration incomplete. Fix issues above before running benchmarks.[/yellow]")
+        return False
+
+
+async def auth_verify(target_url: Optional[str] = None):
+    """Verify authentication by actually signing in."""
+    console.print("\n[bold]Authentication Verification[/bold]\n")
+    
+    try:
+        with console.status("Authenticating..."):
+            client, auth_result = await ensure_admin_authenticated(
+                base_url=target_url,
+                wait_for_service=True,
+                service_wait_retries=10,
+            )
+        
+        console.print(f"  [green]✓ Authentication successful![/green]")
+        console.print(f"  User ID: [cyan]{auth_result.user.id}[/cyan]")
+        console.print(f"  Email: [cyan]{auth_result.user.email}[/cyan]")
+        console.print(f"  Role: [cyan]{auth_result.user.role}[/cyan]")
+        
+        if auth_result.is_new_signup:
+            console.print(f"  [yellow]Note: Created new admin account (first run)[/yellow]")
+        
+        # Clean up
+        await client.close()
+        
+        console.print("\n[green]Ready to run benchmarks![/green]")
+        return True
+        
+    except ServiceNotReadyError as e:
+        console.print(f"  [red]✗ Service not ready: {e}[/red]")
+        return False
+    except AuthenticationError as e:
+        console.print(f"  [red]✗ Authentication failed: {e}[/red]")
+        return False
+    except Exception as e:
+        console.print(f"  [red]✗ Unexpected error: {e}[/red]")
+        return False
+
+
 def list_benchmarks():
     """List available benchmarks."""
     console.print("\n[bold]Available Benchmarks:[/bold]\n")
@@ -187,6 +263,33 @@ def main():
         help="List available benchmarks",
     )
     
+    # Auth command with subcommands
+    auth_parser = subparsers.add_parser(
+        "auth",
+        help="Authentication commands",
+    )
+    auth_subparsers = auth_parser.add_subparsers(dest="auth_command", help="Auth subcommands")
+    
+    # Auth check subcommand
+    auth_check_parser = auth_subparsers.add_parser(
+        "check",
+        help="Check auth configuration and service status (no actual auth)",
+    )
+    auth_check_parser.add_argument(
+        "-u", "--url",
+        help="Target Open WebUI URL (default: from config)",
+    )
+    
+    # Auth verify subcommand
+    auth_verify_parser = auth_subparsers.add_parser(
+        "verify",
+        help="Verify authentication by signing in",
+    )
+    auth_verify_parser.add_argument(
+        "-u", "--url",
+        help="Target Open WebUI URL (default: from config)",
+    )
+    
     # Run command
     run_parser = subparsers.add_parser(
         "run",
@@ -229,6 +332,15 @@ def main():
         list_profiles()
     elif args.command == "list":
         list_benchmarks()
+    elif args.command == "auth":
+        if args.auth_command == "check":
+            success = asyncio.run(auth_check(target_url=args.url))
+            sys.exit(0 if success else 1)
+        elif args.auth_command == "verify":
+            success = asyncio.run(auth_verify(target_url=args.url))
+            sys.exit(0 if success else 1)
+        else:
+            auth_parser.print_help()
     elif args.command == "run":
         try:
             if args.benchmark == "all":

benchmark/scenarios/channels.py

@@ -22,6 +22,7 @@ from benchmark.core.config import BenchmarkConfig
 from benchmark.core.metrics import MetricsCollector, BenchmarkResult
 from benchmark.clients.http_client import OpenWebUIClient, ClientPool
 from benchmark.clients.websocket_client import WebSocketClient, WebSocketPool
+from benchmark.auth import ensure_admin_authenticated, AuthenticationError, ServiceNotReadyError
 
 
 console = Console()
@@ -58,11 +59,24 @@ class ChannelConcurrencyBenchmark(BaseBenchmark):
     description = "Test concurrent user capacity in Open WebUI Channels"
     version = "1.0.0"
     
-    def __init__(self, config: BenchmarkConfig):
-        """Initialize the channel concurrency benchmark."""
+    def __init__(
+        self,
+        config: BenchmarkConfig,
+        admin_client: Optional[OpenWebUIClient] = None,
+    ):
+        """
+        Initialize the channel concurrency benchmark.
+        
+        Args:
+            config: Benchmark configuration
+            admin_client: Optional pre-authenticated admin client. If provided,
+                         skips authentication during setup. If None, performs
+                         authentication using credentials from environment.
+        """
         super().__init__(config)
         
-        self._admin_client: Optional[OpenWebUIClient] = None
+        self._admin_client: Optional[OpenWebUIClient] = admin_client
+        self._owns_admin_client: bool = admin_client is None  # Track if we should close it
         self._client_pool: Optional[ClientPool] = None
         self._ws_pool: Optional[WebSocketPool] = None
         self._test_channel_id: Optional[str] = None
@@ -72,45 +86,28 @@ class ChannelConcurrencyBenchmark(BaseBenchmark):
         """
         Set up the benchmark environment.
         
-        Uses admin credentials to:
-        1. Create a test channel
-        2. Create temporary benchmark users dynamically
+        If no pre-authenticated admin client was provided, uses the auth
+        entrypoint to authenticate. Then:
+        1. Creates a test channel
+        2. Initializes client pool for benchmark users
         
-        Requires ADMIN_USER_EMAIL and ADMIN_USER_PASSWORD in environment.
+        Requires ADMIN_USER_EMAIL and ADMIN_USER_PASSWORD in environment
+        if admin_client was not provided to __init__.
         """
-        # Validate that we have admin credentials configured
-        if not self.config.admin_user:
-            raise RuntimeError(
-                "Admin user credentials not configured. "
-                "Set ADMIN_USER_EMAIL and ADMIN_USER_PASSWORD environment variables."
-            )
-        
-        # Create admin client
-        self._admin_client = OpenWebUIClient(
-            self.config.target_url,
-            self.config.request_timeout,
-        )
-        await self._admin_client.connect()
-        
-        # Wait for service to be ready
-        if not await self._admin_client.wait_for_ready():
-            raise RuntimeError("Open WebUI service not ready")
-        
-        # Authenticate admin
-        admin_config = self.config.admin_user
-        try:
-            await self._admin_client.signin(admin_config.email, admin_config.password)
-        except Exception:
+        # Use provided client or authenticate via entrypoint
+        if self._admin_client is None:
             try:
-                await self._admin_client.signup(
-                    admin_config.email,
-                    admin_config.password,
-                    admin_config.name,
-                )
-            except Exception as e:
-                raise RuntimeError(
-                    f"Failed to authenticate admin ({admin_config.email}): {e}"
+                self._admin_client, _auth_result = await ensure_admin_authenticated(
+                    base_url=self.config.target_url,
+                    timeout=self.config.request_timeout,
+                    wait_for_service=True,
                 )
+                # Note: We don't print here to avoid interfering with progress display
+                # Use `owb auth verify` beforehand to see auth details
+            except ServiceNotReadyError as e:
+                raise RuntimeError(f"Open WebUI service not ready: {e}")
+            except AuthenticationError as e:
+                raise RuntimeError(f"Admin authentication failed: {e}")
         
         # Create test channel
         channel_name = f"benchmark-channel-{int(time.time())}"
@@ -462,8 +459,8 @@ class ChannelConcurrencyBenchmark(BaseBenchmark):
             except Exception:
                 pass
         
-        # Close admin client
-        if self._admin_client:
+        # Close admin client only if we created it
+        if self._admin_client and self._owns_admin_client:
             await self._admin_client.close()
         
         # Close any remaining pooled clients
@@ -487,48 +484,49 @@ class ChannelWebSocketBenchmark(BaseBenchmark):
     description = "Test WebSocket scalability in Open WebUI Channels"
     version = "1.0.0"
     
-    def __init__(self, config: BenchmarkConfig):
-        """Initialize the WebSocket benchmark."""
+    def __init__(
+        self,
+        config: BenchmarkConfig,
+        admin_client: Optional[OpenWebUIClient] = None,
+    ):
+        """
+        Initialize the WebSocket benchmark.
+        
+        Args:
+            config: Benchmark configuration
+            admin_client: Optional pre-authenticated admin client. If provided,
+                         skips authentication during setup. If None, performs
+                         authentication using credentials from environment.
+        """
         super().__init__(config)
         
-        self._admin_client: Optional[OpenWebUIClient] = None
+        self._admin_client: Optional[OpenWebUIClient] = admin_client
+        self._owns_admin_client: bool = admin_client is None  # Track if we should close it
         self._client_pool: Optional[ClientPool] = None
         self._ws_pool: Optional[WebSocketPool] = None
         self._test_channel_id: Optional[str] = None
     
     async def setup(self) -> None:
-        """Set up benchmark environment."""
-        # Similar to ChannelConcurrencyBenchmark setup
-        self._admin_client = OpenWebUIClient(
-            self.config.target_url,
-            self.config.request_timeout,
-        )
-        await self._admin_client.connect()
+        """
+        Set up benchmark environment.
         
-        if not await self._admin_client.wait_for_ready():
-            raise RuntimeError("Open WebUI service not ready")
-        
-        # Authenticate admin
-        admin_config = self.config.admin_user
-        if not admin_config:
-            raise RuntimeError(
-                "Admin credentials not configured. "
-                "Set ADMIN_USER_EMAIL and ADMIN_USER_PASSWORD."
-            )
-        
-        try:
-            await self._admin_client.signin(admin_config.email, admin_config.password)
-        except Exception:
+        If no pre-authenticated admin client was provided, uses the auth
+        entrypoint to authenticate. Then creates test channel and initializes pools.
+        """
+        # Use provided client or authenticate via entrypoint
+        if self._admin_client is None:
             try:
-                await self._admin_client.signup(
-                    admin_config.email,
-                    admin_config.password,
-                    admin_config.name,
-                )
-            except Exception as e:
-                raise RuntimeError(
-                    f"Failed to authenticate admin ({admin_config.email}): {e}"
+                self._admin_client, _auth_result = await ensure_admin_authenticated(
+                    base_url=self.config.target_url,
+                    timeout=self.config.request_timeout,
+                    wait_for_service=True,
                 )
+                # Note: We don't print here to avoid interfering with progress display
+                # Use `owb auth verify` beforehand to see auth details
+            except ServiceNotReadyError as e:
+                raise RuntimeError(f"Open WebUI service not ready: {e}")
+            except AuthenticationError as e:
+                raise RuntimeError(f"Admin authentication failed: {e}")
         
         # Create test channel
         channel_name = f"benchmark-ws-channel-{int(time.time())}"
@@ -621,7 +619,8 @@ class ChannelWebSocketBenchmark(BaseBenchmark):
             except Exception:
                 pass
         
-        if self._admin_client:
+        # Close admin client only if we created it
+        if self._admin_client and self._owns_admin_client:
             await self._admin_client.close()
         
         if self._client_pool:

examples/custom_benchmark.py

@@ -2,7 +2,8 @@
 Example showing how to create a custom benchmark.
 
 This demonstrates extending the base benchmark class to create
-new benchmark scenarios.
+new benchmark scenarios, using the auth entrypoint for consistent
+authentication handling.
 """
 
 import asyncio
@@ -13,6 +14,11 @@ from benchmark.core.base import BaseBenchmark
 from benchmark.core.config import BenchmarkConfig
 from benchmark.core.metrics import BenchmarkResult
 from benchmark.clients.http_client import OpenWebUIClient
+from benchmark.auth import (
+    ensure_admin_authenticated,
+    AuthenticationError,
+    ServiceNotReadyError,
+)
 
 
 class APILatencyBenchmark(BaseBenchmark):
@@ -34,35 +40,37 @@ class APILatencyBenchmark(BaseBenchmark):
         ("GET", "/health"),
     ]
     
-    def __init__(self, config: BenchmarkConfig):
-        """Initialize the benchmark."""
+    def __init__(
+        self,
+        config: BenchmarkConfig,
+        admin_client: Optional[OpenWebUIClient] = None,
+    ):
+        """
+        Initialize the benchmark.
+        
+        Args:
+            config: Benchmark configuration
+            admin_client: Optional pre-authenticated admin client
+        """
         super().__init__(config)
-        self._client: Optional[OpenWebUIClient] = None
+        self._client: Optional[OpenWebUIClient] = admin_client
+        self._owns_client: bool = admin_client is None
     
     async def setup(self) -> None:
-        """Set up the benchmark environment."""
-        self._client = OpenWebUIClient(
-            self.config.target_url,
-            self.config.request_timeout,
-        )
-        await self._client.connect()
-        
-        # Wait for service
-        if not await self._client.wait_for_ready():
-            raise RuntimeError("Open WebUI service not ready")
-        
-        # Authenticate
-        try:
-            await self._client.signin(
-                "admin@benchmark.local",
-                "benchmark_admin_123",
-            )
-        except Exception:
-            await self._client.signup(
-                "admin@benchmark.local",
-                "benchmark_admin_123",
-                "Benchmark Admin",
-            )
+        """Set up the benchmark environment using auth entrypoint."""
+        if self._client is None:
+            try:
+                self._client, auth_result = await ensure_admin_authenticated(
+                    base_url=self.config.target_url,
+                    timeout=self.config.request_timeout,
+                )
+                print(f"Authenticated as: {auth_result.user.email}")
+                if auth_result.is_new_signup:
+                    print("(Created new admin account)")
+            except ServiceNotReadyError as e:
+                raise RuntimeError(f"Service not ready: {e}")
+            except AuthenticationError as e:
+                raise RuntimeError(f"Authentication failed: {e}")
     
     async def run(self) -> BenchmarkResult:
         """Execute the benchmark."""
@@ -112,7 +120,7 @@ class APILatencyBenchmark(BaseBenchmark):
     
     async def teardown(self) -> None:
         """Clean up resources."""
-        if self._client:
+        if self._client and self._owns_client:
             await self._client.close()
 
 

examples/run_channel_benchmark.py

@@ -1,7 +1,14 @@
 """
 Example script demonstrating how to run the channel concurrency benchmark.
 
-This script can be run directly or used as a reference for programmatic benchmark execution.
+This script shows two approaches:
+1. Let the benchmark handle authentication internally (simple)
+2. Pre-authenticate using the entrypoint and pass the client (explicit)
+
+The explicit approach is useful when you want to:
+- Validate authentication before starting benchmarks
+- Share one authenticated session across multiple benchmarks
+- Handle auth errors separately from benchmark errors
 """
 
 import asyncio
@@ -10,11 +17,16 @@ from pathlib import Path
 from benchmark.core.config import load_config
 from benchmark.core.runner import BenchmarkRunner
 from benchmark.scenarios.channels import ChannelConcurrencyBenchmark
+from benchmark.auth import ensure_admin_authenticated
 
 
-async def main():
-    """Run the channel concurrency benchmark with custom settings."""
+async def main_simple():
+    """
+    Run benchmark with internal authentication (simple approach).
     
+    The benchmark will handle authentication using credentials from
+    environment variables (ADMIN_USER_EMAIL, ADMIN_USER_PASSWORD).
+    """
     # Load configuration with the default compute profile
     config = load_config(
         profile_id="default",
@@ -24,10 +36,10 @@ async def main():
     )
     
     # Customize channel benchmark settings
-    config.channels.max_concurrent_users = 50  # Test up to 50 users
-    config.channels.user_step_size = 10        # Increase by 10 users each level
-    config.channels.sustain_time = 20          # Run each level for 20 seconds
-    config.channels.message_frequency = 0.5    # 0.5 messages per second per user
+    config.channels.max_concurrent_users = 50
+    config.channels.user_step_size = 10
+    config.channels.sustain_time = 20
+    config.channels.message_frequency = 0.5
     
     # Create benchmark runner
     runner = BenchmarkRunner(
@@ -36,7 +48,7 @@ async def main():
         output_dir=Path("./results"),
     )
     
-    # Run the benchmark
+    # Run the benchmark (handles auth internally)
     print("Starting Channel Concurrency Benchmark...")
     print(f"Target: {config.target_url}")
     print(f"Max users: {config.channels.max_concurrent_users}")
@@ -47,16 +59,77 @@ async def main():
     # Display results
     runner.display_final_summary()
     
-    # Access specific metrics
-    print(f"\n--- Benchmark Complete ---")
-    print(f"Maximum sustainable users: {result.metadata.get('max_sustainable_users', 'N/A')}")
-    print(f"Average response time: {result.avg_response_time_ms:.2f}ms")
-    print(f"P95 response time: {result.p95_response_time_ms:.2f}ms")
-    print(f"Error rate: {result.error_rate_percent:.2f}%")
-    print(f"Passed: {'Yes' if result.passed else 'No'}")
-    
     return result
 
 
+async def main_explicit():
+    """
+    Run benchmark with explicit pre-authentication (recommended approach).
+    
+    This approach uses the auth entrypoint to authenticate before
+    creating the benchmark, giving you more control over error handling.
+    """
+    # Load configuration
+    config = load_config(
+        profile_id="default",
+        overrides={
+            "target_url": "http://localhost:3000",
+        }
+    )
+    
+    config.channels.max_concurrent_users = 50
+    config.channels.user_step_size = 10
+    config.channels.sustain_time = 20
+    config.channels.message_frequency = 0.5
+    
+    # Pre-authenticate using the entrypoint
+    print("Authenticating...")
+    client, auth_result = await ensure_admin_authenticated(
+        base_url=config.target_url,
+        timeout=config.request_timeout,
+    )
+    
+    print(f"Authenticated as: {auth_result.user.email} (role: {auth_result.user.role})")
+    if auth_result.is_new_signup:
+        print("Note: Created new admin account (first run)")
+    print()
+    
+    try:
+        # Create benchmark runner
+        runner = BenchmarkRunner(
+            config=config,
+            profile_id="default",
+            output_dir=Path("./results"),
+        )
+        
+        # Create benchmark with pre-authenticated client
+        benchmark = ChannelConcurrencyBenchmark(config, admin_client=client)
+        
+        # Run the benchmark
+        print("Starting Channel Concurrency Benchmark...")
+        print(f"Target: {config.target_url}")
+        print(f"Max users: {config.channels.max_concurrent_users}")
+        print()
+        
+        result = await runner.run_benchmark(benchmark)
+        
+        # Display results
+        runner.display_final_summary()
+        
+        # Access specific metrics
+        print(f"\n--- Benchmark Complete ---")
+        print(f"Maximum sustainable users: {result.metadata.get('max_sustainable_users', 'N/A')}")
+        print(f"Average response time: {result.avg_response_time_ms:.2f}ms")
+        print(f"P95 response time: {result.p95_response_time_ms:.2f}ms")
+        print(f"Error rate: {result.error_rate_percent:.2f}%")
+        print(f"Passed: {'Yes' if result.passed else 'No'}")
+        
+        return result
+    finally:
+        # Clean up the client we created
+        await client.close()
+
+
 if __name__ == "__main__":
-    asyncio.run(main())
+    # Use the explicit approach by default
+    asyncio.run(main_explicit())