Spaces:

techprotrade
/

annator-command-center

Paused

annator-command-center / tests /integration /workflows /conftest.py

Deploy ATOM FastAPI command center runtime (part 8)

aef804e verified 3 months ago

11.7 kB

"""
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]
    }

	"""
	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]
	}