diodiogod-TTS-Audio-Suite/tests/conftest.py

"""
TTS-Audio-Suite Test Configuration

Session-scoped fixtures for ComfyUI server management and API testing.
"""

# ============================================================================
# CRITICAL: Mock ComfyUI modules BEFORE anything else
# This prevents import errors when pytest discovers and imports test modules
# ============================================================================
import os
import sys
import tempfile
import types
from unittest.mock import MagicMock

# Set testing environment
os.environ['COMFYUI_TESTING'] = '1'
os.environ['PYTEST_CURRENT_TEST'] = 'true'

# Mock ComfyUI modules at module level - BEFORE any imports can trigger __init__.py
_MOCK_MODULES = [
    'comfy',
    'comfy.model_management',
    'comfy.utils',
    'nodes',
    'server',
    'execution',
    'comfy_extras',
]

for _module_name in _MOCK_MODULES:
    if _module_name not in sys.modules:
        sys.modules[_module_name] = MagicMock()

# folder_paths cannot be a raw MagicMock because code under test uses
# folder_paths.models_dir in os.path.join/os.makedirs calls. A MagicMock there
# creates junk folders like ./MagicMock/mock.models_dir/... on disk.
if 'folder_paths' not in sys.modules:
    mock_folder_paths = types.ModuleType("folder_paths")
    mock_models_dir = tempfile.mkdtemp(prefix="tts_suite_test_models_")
    mock_input_dir = tempfile.mkdtemp(prefix="tts_suite_test_input_")
    mock_output_dir = tempfile.mkdtemp(prefix="tts_suite_test_output_")
    mock_temp_dir = tempfile.mkdtemp(prefix="tts_suite_test_temp_")

    mock_folder_paths.models_dir = mock_models_dir
    mock_folder_paths.input_directory = mock_input_dir
    mock_folder_paths.output_directory = mock_output_dir
    mock_folder_paths.temp_directory = mock_temp_dir

    mock_folder_paths.get_folder_paths = lambda *args, **kwargs: []
    mock_folder_paths.add_model_folder_path = lambda *args, **kwargs: None
    mock_folder_paths.get_input_directory = lambda: mock_input_dir
    mock_folder_paths.get_output_directory = lambda: mock_output_dir
    mock_folder_paths.get_temp_directory = lambda: mock_temp_dir
    mock_folder_paths.get_annotated_filepath = lambda path: path
    mock_folder_paths.exists_annotated_filepath = lambda path: os.path.exists(path)

    sys.modules['folder_paths'] = mock_folder_paths

# Now safe to import pytest and other modules
import pytest
import subprocess
import time
import json
import requests
from pathlib import Path
from typing import Dict, Any, Optional

# Tell pytest to ignore certain files during collection
collect_ignore = ["__init__.py", "nodes.py"]

# Path configuration
CUSTOM_NODE_ROOT = Path(__file__).parent.parent  # tests/ -> TTS-Audio-Suite/
DEFAULT_COMFY_ROOT = CUSTOM_NODE_ROOT.parent.parent  # Navigate to Comfy-new
COMFY_ROOT = Path(os.environ.get("TTS_SUITE_TEST_COMFY_ROOT", str(DEFAULT_COMFY_ROOT)))

if os.name == "nt":
    DEFAULT_VENV_PYTHON = COMFY_ROOT / "venv" / "Scripts" / "python.exe"
else:
    DEFAULT_VENV_PYTHON = COMFY_ROOT / "venv" / "bin" / "python"

VENV_PYTHON = Path(os.environ.get("TTS_SUITE_TEST_VENV_PYTHON", str(DEFAULT_VENV_PYTHON)))


class ComfyUIAPIClient:
    """Helper class for interacting with ComfyUI API during tests"""
    
    def __init__(self, base_url: str):
        self.base_url = base_url
        self.client_id = "pytest-test-client"
    
    def get_system_stats(self) -> Dict[str, Any]:
        """Get system stats to verify server is running"""
        response = requests.get(f"{self.base_url}/system_stats", timeout=5)
        response.raise_for_status()
        return response.json()
    
    def get_object_info(self) -> Dict[str, Any]:
        """Get all registered node info - useful for verifying nodes loaded"""
        response = requests.get(f"{self.base_url}/object_info", timeout=30)
        response.raise_for_status()
        return response.json()
    
    def queue_prompt(self, workflow: Dict[str, Any]) -> Dict[str, Any]:
        """Queue a workflow for execution"""
        response = requests.post(
            f"{self.base_url}/prompt",
            json={"prompt": workflow, "client_id": self.client_id},
            timeout=30
        )
        response.raise_for_status()
        return response.json()
    
    def get_history(self, prompt_id: str) -> Dict[str, Any]:
        """Get execution history for a prompt"""
        response = requests.get(
            f"{self.base_url}/history/{prompt_id}",
            timeout=10
        )
        response.raise_for_status()
        return response.json()
    
    def wait_for_completion(self, prompt_id: str, timeout: int = 120) -> Dict[str, Any]:
        """Wait for workflow to complete"""
        start = time.time()
        while time.time() - start < timeout:
            try:
                history = self.get_history(prompt_id)
                if prompt_id in history:
                    outputs = history[prompt_id].get("outputs", {})
                    status = history[prompt_id].get("status", {})
                    
                    # Check if completed
                    if status.get("completed", False):
                        return history[prompt_id]
                    
                    # Check for error
                    status_str = status.get("status_str", "")
                    if status_str == "error":
                        raise RuntimeError(
                            f"Workflow execution failed: {status}"
                        )
            except requests.RequestException:
                pass  # Server might be busy
            
            time.sleep(0.5)
        
        raise TimeoutError(
            f"Workflow did not complete within {timeout}s (prompt_id: {prompt_id})"
        )
    
    def execute_workflow(self, workflow: Dict[str, Any], timeout: int = 120) -> Dict[str, Any]:
        """Queue workflow and wait for completion"""
        result = self.queue_prompt(workflow)
        prompt_id = result["prompt_id"]
        return self.wait_for_completion(prompt_id, timeout)
    
    def node_exists(self, node_class: str) -> bool:
        """Check if a node class is registered"""
        try:
            object_info = self.get_object_info()
            return node_class in object_info
        except Exception:
            return False


@pytest.fixture(scope="session")
def comfyui_server():
    """
    Start ComfyUI server for the entire test session.
    Automatically shuts down when tests complete.
    
    Yields:
        dict with 'url' and 'process' keys
    """
    print("\n🚀 Starting ComfyUI server for integration tests...")
    
    # Check if server is already running
    comfy_host = os.environ.get("TTS_SUITE_TEST_HOST", "127.0.0.1")
    comfy_port = int(os.environ.get("TTS_SUITE_TEST_PORT", "8188"))
    server_url = f"http://{comfy_host}:{comfy_port}"
    try:
        response = requests.get(f"{server_url}/system_stats", timeout=2)
        if response.status_code == 200:
            print(f"✅ ComfyUI server already running at {server_url}")
            yield {"url": server_url, "process": None, "external": True}
            return
    except (requests.ConnectionError, requests.Timeout):
        pass  # Server not running, we'll start it
    
    # Start ComfyUI in subprocess with process group isolation and a clean env (no COMFYUI_TESTING)
    creation_flags = subprocess.CREATE_NEW_PROCESS_GROUP if sys.platform == "win32" else 0
    server_env = os.environ.copy()
    server_env.pop("COMFYUI_TESTING", None)
    server_env.pop("PYTEST_CURRENT_TEST", None)

    log_path = CUSTOM_NODE_ROOT / "tests" / "comfyui_server.log"
    log_file = open(log_path, "w", encoding="utf-8", errors="replace")

    process = subprocess.Popen(
        [
            str(VENV_PYTHON),
            "main.py",
            "--listen", comfy_host,
            "--port", str(comfy_port),
            "--disable-auto-launch",
            "--cpu"  # Use CPU for faster startup in tests
        ],
        cwd=str(COMFY_ROOT),
        stdout=log_file,
        stderr=subprocess.STDOUT,
        creationflags=creation_flags,
        preexec_fn=os.setsid if sys.platform != "win32" else None,
        env=server_env
    )
    
    # Wait for server to be ready
    max_retries = 120  # 2 minutes timeout for model loading
    retry_count = 0
    
    while retry_count < max_retries:
        try:
            response = requests.get(f"{server_url}/system_stats", timeout=2)
            if response.status_code == 200:
                print(f"✅ ComfyUI server ready at {server_url}")
                break
        except (requests.ConnectionError, requests.Timeout):
            time.sleep(1)
            retry_count += 1
            if retry_count % 15 == 0:
                print(f"⏳ Waiting for server... ({retry_count}s)")
    
    if retry_count >= max_retries:
        # Try to get any error output
        process.terminate()
        try:
            log_file.close()
            with open(log_path, "r", encoding="utf-8", errors="replace") as f:
                error_lines = f.read()[-2000:]
        except Exception:
            error_lines = "Could not capture output"
        
        raise RuntimeError(
            f"❌ ComfyUI server failed to start within {max_retries}s\n"
            f"Last output:\n{error_lines}"
        )
    
    try:
        yield {"url": server_url, "process": process, "external": False}
    finally:
        # Teardown
        if process and process.poll() is None:
            print("\n🛑 Shutting down ComfyUI server...")
            if sys.platform == "win32":
                process.terminate()
            else:
                try:
                    import signal
                    os.killpg(os.getpgid(process.pid), signal.SIGTERM)
                except ProcessLookupError:
                    process.terminate()
            try:
                process.wait(timeout=10)
            except subprocess.TimeoutExpired:
                process.kill()
                process.wait()
            print("✅ Server stopped")
        log_file.close()


@pytest.fixture
def api_client(comfyui_server) -> ComfyUIAPIClient:
    """Get API client connected to running ComfyUI server"""
    return ComfyUIAPIClient(comfyui_server["url"])


@pytest.fixture
def workflow_fixtures_path() -> Path:
    """Path to workflow test fixtures"""
    return CUSTOM_NODE_ROOT / "tests" / "fixtures" / "workflows"


@pytest.fixture
def sample_voice_path() -> Optional[Path]:
    """Path to a sample voice file for testing (if available)"""
    voices_dir = CUSTOM_NODE_ROOT / "voices_examples"
    if voices_dir.exists():
        # Find any .wav file
        for wav_file in voices_dir.rglob("*.wav"):
            return wav_file
    return None


# ============================================================================
# Unit test fixtures (no server required)
# ============================================================================

@pytest.fixture
def sample_srt_content() -> str:
    """Sample SRT content for parser testing"""
    return """1
00:00:01,000 --> 00:00:04,000
Hello, this is the first subtitle.

2
00:00:05,000 --> 00:00:08,500
And this is the second one.

3
00:00:10,000 --> 00:00:12,000
Final subtitle here.
"""


@pytest.fixture
def sample_text_with_pauses() -> str:
    """Sample text with pause tags for testing"""
    return "Hello! [pause:2] How are you today? [wait:500ms] I hope you're doing well."


@pytest.fixture
def sample_text_with_characters() -> str:
    """Sample text with character tags for testing"""
    return """[Alice]
Hello everyone! Welcome to the show.

[Bob]
Thanks Alice! Great to be here.

[Alice]
Let's get started!
"""