""" PyTest configuration and fixtures for API tests. """ import os import sys import json import pytest import requests import tempfile import time import yaml from pathlib import Path from typing import Dict, Generator, Optional, Tuple # Import test utilities import sys import os from pathlib import Path # Get the absolute path to the current file (conftest.py) current_file = Path(os.path.abspath(__file__)) # Get the directory containing the current file (the tests-api directory) tests_api_dir = current_file.parent # Add the tests-api directory to the Python path if str(tests_api_dir) not in sys.path: sys.path.insert(0, str(tests_api_dir)) # Apply mocks for ComfyUI imports from mocks.patch import apply_mocks apply_mocks() # Now we can import from utils.validation from utils.validation import load_openapi_spec # Default test configuration DEFAULT_TEST_CONFIG = { "server_url": "http://localhost:8188", "server_timeout": 2, # seconds "wait_between_requests": 0.5, # seconds "max_retries": 3, } @pytest.fixture(scope="session") def test_config() -> Dict: """ Load test configuration from environment variables or use defaults. """ config = DEFAULT_TEST_CONFIG.copy() # Override from environment variables if present if "COMFYUI_SERVER_URL" in os.environ: config["server_url"] = os.environ["COMFYUI_SERVER_URL"] return config @pytest.fixture(scope="session") def server_url(test_config: Dict) -> str: """ Get the server URL from the test configuration. """ return test_config["server_url"] @pytest.fixture(scope="session") def openapi_spec() -> Dict: """ Load the OpenAPI specification. """ return load_openapi_spec() @pytest.fixture(scope="session") def api_client(server_url: str, test_config: Dict) -> requests.Session: """ Create a requests Session for API calls. """ session = requests.Session() # Check if the server is running try: response = session.get(f"{server_url}/", timeout=test_config["server_timeout"]) response.raise_for_status() except (requests.ConnectionError, requests.Timeout, requests.HTTPError): pytest.skip("ComfyUI server is not running or not accessible") return session @pytest.fixture(scope="function") def temp_dir() -> Generator[Path, None, None]: """ Create a temporary directory for test files. """ with tempfile.TemporaryDirectory() as temp_dir: yield Path(temp_dir) class SecurityLevelContext: """ Context manager for setting and restoring security levels. """ def __init__(self, api_client: requests.Session, server_url: str, security_level: str): self.api_client = api_client self.server_url = server_url self.security_level = security_level self.original_level = None async def __aenter__(self): # Get the current security level (not directly exposed in API, would require more setup) # For now, we'll just set the new level # Set the new security level # Note: In a real implementation, we would need a way to set this # This is a placeholder - the actual implementation would depend on how # security levels are managed in ComfyUI-Manager return self async def __aexit__(self, exc_type, exc_val, exc_tb): # Restore the original security level if needed pass @pytest.fixture def security_level_context(api_client: requests.Session, server_url: str): """ Create a context manager for setting security levels. """ return lambda level: SecurityLevelContext(api_client, server_url, level) def make_api_url(server_url: str, path: str) -> str: """ Construct a full API URL from the server URL and path. """ # Ensure the path starts with a slash if not path.startswith("/"): path = f"/{path}" # Remove trailing slash from server_url if present if server_url.endswith("/"): server_url = server_url[:-1] return f"{server_url}{path}" @pytest.fixture def api_request(api_client: requests.Session, server_url: str, test_config: Dict): """ Helper function for making API requests with consistent behavior. """ def _request( method: str, path: str, params: Optional[Dict] = None, json_data: Optional[Dict] = None, headers: Optional[Dict] = None, expected_status: int = 200, retry_on_error: bool = True, ) -> Tuple[requests.Response, Optional[Dict]]: """ Make an API request with automatic validation. Args: method: HTTP method path: API path params: Query parameters json_data: JSON request body headers: HTTP headers expected_status: Expected HTTP status code retry_on_error: Whether to retry on connection errors Returns: Tuple of (Response object, JSON response data or None) """ method = method.lower() url = make_api_url(server_url, path) if headers is None: headers = {} # Add common headers headers.setdefault("Accept", "application/json") # Sleep between requests to avoid overwhelming the server time.sleep(test_config["wait_between_requests"]) retries = test_config["max_retries"] if retry_on_error else 0 last_exception = None for attempt in range(retries + 1): try: if method == "get": response = api_client.get(url, params=params, headers=headers) elif method == "post": response = api_client.post(url, params=params, json=json_data, headers=headers) elif method == "put": response = api_client.put(url, params=params, json=json_data, headers=headers) elif method == "delete": response = api_client.delete(url, params=params, headers=headers) else: raise ValueError(f"Unsupported HTTP method: {method}") # Check status code assert response.status_code == expected_status, ( f"Expected status code {expected_status}, got {response.status_code}" ) # Parse JSON response if possible json_response = None if response.headers.get("Content-Type", "").startswith("application/json"): try: json_response = response.json() except json.JSONDecodeError: if expected_status == 200: raise ValueError("Response was not valid JSON") return response, json_response except (requests.ConnectionError, requests.Timeout) as e: last_exception = e if attempt < retries: # Wait before retrying time.sleep(1) continue break if last_exception: raise last_exception raise RuntimeError("Failed to make API request") return _request