tests/integration/workflows/conftest.py

"""
Shared fixtures for workflow integration tests.

Provides database sessions, mock LLM providers, mock WebSocket managers,
and sample workflow/episode data for integration testing.

Key principles:
- Real database (SQLite in-memory) for persistence testing
- Mocked external services (LLM providers, WebSocket) to prevent flakiness
- Transaction rollback after each test for isolation
"""

import os
import pytest
import tempfile
import uuid
from datetime import datetime, timedelta
from unittest.mock import AsyncMock, MagicMock
from sqlalchemy import create_engine
from sqlalchemy.orm import Session, sessionmaker

# Set TESTING environment variable BEFORE any imports
os.environ["TESTING"] = "1"

# Add parent directory to path for imports
import sys
from pathlib import Path
sys.path.insert(0, str(Path(__file__).parent.parent.parent.parent))

from core.database import Base
from core.models import (
    WorkflowExecution, WorkflowExecutionStatus,
    AgentEpisode, EpisodeSegment, EpisodeOutcome,
    ChatSession, CanvasAudit, AgentFeedback
)


@pytest.fixture(scope="function")
def db_session():
    """
    Create a fresh in-memory database for each test.

    Uses SQLite in-memory database with transaction rollback for test isolation.
    All tables are created before the test and cleaned up after.
    """
    # Use in-memory SQLite for tests
    engine = create_engine(
        "sqlite:///:memory:",
        connect_args={"check_same_thread": False},
        echo=False
    )

    # Create all tables
    try:
        Base.metadata.create_all(engine, checkfirst=True)
    except Exception as e:
        # If create_all fails, create tables individually
        for table in Base.metadata.tables.values():
            try:
                table.create(engine, checkfirst=True)
            except Exception:
                continue

    # Create session
    TestingSessionLocal = sessionmaker(autocommit=False, autoflush=False, bind=engine)
    session = TestingSessionLocal()

    yield session

    # Cleanup
    session.close()
    engine.dispose()


@pytest.fixture(scope="function")
def mock_llm_provider():
    """
    AsyncMock for LLM provider calls.

    Mocks both call_llm() and stream_llm() methods to prevent
    external API calls during testing.
    """
    mock_provider = AsyncMock()

    # Mock call_llm to return success response
    async def mock_call_llm(*args, **kwargs):
        return {
            "result": "success",
            "content": "Mock LLM response",
            "usage": {
                "prompt_tokens": 10,
                "completion_tokens": 20,
                "total_tokens": 30
            }
        }

    # Mock stream_llm to return async generator
    async def mock_stream_llm(*args, **kwargs):
        chunks = [
            {"choices": [{"delta": {"content": "Mock "}, "finish_reason": None}], "usage": None},
            {"choices": [{"delta": {"content": "streaming "}, "finish_reason": None}], "usage": None},
            {"choices": [{"delta": {"content": "response"}, "finish_reason": None}], "usage": None},
            {"choices": [{"delta": {}, "finish_reason": "stop"}],
             "usage": {"prompt_tokens": 10, "completion_tokens": 20, "total_tokens": 30}}
        ]
        for chunk in chunks:
            yield chunk

    mock_provider.call_llm = mock_call_llm
    mock_provider.stream_llm = mock_stream_llm

    return mock_provider


@pytest.fixture(scope="function")
def mock_websocket_manager():
    """
    MagicMock for WebSocket notifications.

    Tracks all WebSocket calls for verification in tests.
    """
    mock_ws = MagicMock()

    # Mock notification methods
    mock_ws.notify_workflow_status = AsyncMock()
    mock_ws.broadcast_message = AsyncMock()
    mock_ws.send_message = AsyncMock()

    # Track calls for verification
    mock_ws.reset_mock()

    return mock_ws


@pytest.fixture(scope="function")
def sample_workflow():
    """
    Valid workflow definitions for tests.

    Returns a dict with three workflow variations:
    - simple: 2-step linear execution
    - branching: if/else conditional logic
    - error_handling: try/catch error recovery
    """
    return {
        "simple": {
            "id": "simple_workflow",
            "nodes": [
                {
                    "id": "step1",
                    "title": "First Step",
                    "type": "action",
                    "config": {
                        "action": "collect_data",
                        "service": "data_service",
                        "parameters": {"source": "api"}
                    }
                },
                {
                    "id": "step2",
                    "title": "Second Step",
                    "type": "action",
                    "config": {
                        "action": "process_data",
                        "service": "processing_service",
                        "parameters": {"mode": "batch"}
                    }
                }
            ],
            "connections": [
                {"source": "step1", "target": "step2"}
            ]
        },
        "branching": {
            "id": "branching_workflow",
            "nodes": [
                {
                    "id": "step1",
                    "title": "Evaluate Condition",
                    "type": "action",
                    "config": {
                        "action": "evaluate",
                        "service": "logic_service",
                        "parameters": {}
                    }
                },
                {
                    "id": "true_branch",
                    "title": "True Branch",
                    "type": "action",
                    "config": {
                        "action": "handle_true",
                        "service": "logic_service",
                        "parameters": {"branch": "true"}
                    }
                },
                {
                    "id": "false_branch",
                    "title": "False Branch",
                    "type": "action",
                    "config": {
                        "action": "handle_false",
                        "service": "logic_service",
                        "parameters": {"branch": "false"}
                    }
                }
            ],
            "connections": [
                {"source": "step1", "target": "true_branch", "condition": "${step1.result} == true"},
                {"source": "step1", "target": "false_branch", "condition": "${step1.result} == false"}
            ]
        },
        "error_handling": {
            "id": "error_handling_workflow",
            "nodes": [
                {
                    "id": "step1",
                    "title": "Risky Step",
                    "type": "action",
                    "config": {
                        "action": "risky_operation",
                        "service": "risky_service",
                        "parameters": {"continue_on_error": True}
                    }
                },
                {
                    "id": "step2",
                    "title": "Recovery Step",
                    "type": "action",
                    "config": {
                        "action": "recover",
                        "service": "recovery_service",
                        "parameters": {}
                    }
                }
            ],
            "connections": [
                {"source": "step1", "target": "step2"}
            ]
        }
    }


@pytest.fixture(scope="function")
def sample_episode_context(db_session: Session):
    """
    Episode context for segmentation tests.

    Creates a realistic episode context with:
    - ChatSession with messages
    - Canvas presentation events
    - User feedback data
    """
    # Create ChatSession with messages
    chat_session = ChatSession(
        id=str(uuid.uuid4()),
        user_id="test_user",
        title="Test Episode Session",
        metadata_json={"agent_id": "test_agent"}
    )
    db_session.add(chat_session)

    # Create sample canvas presentation
    canvas_audit = CanvasAudit(
        id=str(uuid.uuid4()),
        canvas_id="test_canvas_123",
        canvas_type="line_chart",
        agent_id="test_agent",
        user_id="test_user",
        session_id=chat_session.id,
        title="Sales Performance Chart",
        config={"data": [1, 2, 3, 4, 5]},
        status="presented"
    )
    db_session.add(canvas_audit)

    # Create sample user feedback
    agent_feedback = AgentFeedback(
        id=str(uuid.uuid4()),
        agent_id="test_agent",
        user_id="test_user",
        session_id=chat_session.id,
        feedback_type="thumbs_up",
        feedback_value=1.0,
        comment="Great analysis!"
    )
    db_session.add(agent_feedback)

    db_session.commit()

    return {
        "chat_session": chat_session,
        "canvas_audit": canvas_audit,
        "agent_feedback": agent_feedback,
        "messages": [
            {
                "role": "user",
                "content": "Show me the sales data",
                "timestamp": datetime.utcnow() - timedelta(minutes=10)
            },
            {
                "role": "assistant",
                "content": "Here's the sales performance chart",
                "timestamp": datetime.utcnow() - timedelta(minutes=9),
                "canvas_id": canvas_audit.canvas_id
            },
            {
                "role": "user",
                "content": "Great analysis!",
                "timestamp": datetime.utcnow() - timedelta(minutes=8),
                "feedback_id": agent_feedback.id
            }
        ]
    }


@pytest.fixture(scope="function")
def sample_episode(db_session: Session):
    """
    Create a sample AgentEpisode with segments for testing.

    Returns an episode with 3 segments representing a simple workflow.
    """
    episode_id = str(uuid.uuid4())

    episode = AgentEpisode(
        id=episode_id,
        agent_id="test_agent",
        user_id="test_user",
        workflow_id="test_workflow",
        outcome=EpisodeOutcome.SUCCESS,
        title="Test Episode",
        summary="Test episode with multiple segments",
        start_time=datetime.utcnow() - timedelta(minutes=10),
        end_time=datetime.utcnow()
    )
    db_session.add(episode)

    # Create segments
    segment1 = EpisodeSegment(
        id=str(uuid.uuid4()),
        episode_id=episode_id,
        segment_type="action",
        title="Data Collection",
        content="Collected data from API",
        start_time=datetime.utcnow() - timedelta(minutes=10),
        end_time=datetime.utcnow() - timedelta(minutes=8),
        metadata={"action": "collect_data", "source": "api"}
    )
    db_session.add(segment1)

    segment2 = EpisodeSegment(
        id=str(uuid.uuid4()),
        episode_id=episode_id,
        segment_type="action",
        title="Data Processing",
        content="Processed data in batch mode",
        start_time=datetime.utcnow() - timedelta(minutes=8),
        end_time=datetime.utcnow() - timedelta(minutes=5),
        metadata={"action": "process_data", "mode": "batch"}
    )
    db_session.add(segment2)

    segment3 = EpisodeSegment(
        id=str(uuid.uuid4()),
        episode_id=episode_id,
        segment_type="result",
        title="Result Presentation",
        content="Presented results to user",
        start_time=datetime.utcnow() - timedelta(minutes=5),
        end_time=datetime.utcnow(),
        metadata={"action": "present_results"}
    )
    db_session.add(segment3)

    db_session.commit()

    return {
        "episode": episode,
        "segments": [segment1, segment2, segment3]
    }