Spaces:
Sleeping
Sleeping
File size: 6,223 Bytes
aef804e | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 195 196 197 198 199 200 201 202 203 204 | """
E2E test fixtures for agent execution workflow tests.
Provides specialized fixtures for end-to-end testing of agent execution
including LLM streaming mocks, WebSocket mocks, and database cleanup.
"""
import os
import pytest
import uuid
from unittest.mock import AsyncMock, MagicMock, patch
from sqlalchemy.orm import Session
from sqlalchemy import text
# 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))
from tests.factories.agent_factory import (
AgentFactory,
StudentAgentFactory,
InternAgentFactory,
SupervisedAgentFactory,
AutonomousAgentFactory
)
from core.models import AgentRegistry, AgentExecution, AgentEpisode, EpisodeSegment
@pytest.fixture(scope="function")
def e2e_db_session(db_session: Session):
"""
E2E database session with aggressive cleanup.
Cleans up all E2E test data after each test to prevent cross-test contamination.
"""
yield db_session
# Aggressive cleanup for E2E tests
try:
# Clean up in order of dependencies
db_session.execute(text("DELETE FROM episode_segments WHERE 1=1"))
db_session.execute(text("DELETE FROM agent_episodes WHERE agent_id LIKE 'test-agent%'"))
db_session.execute(text("DELETE FROM agent_executions WHERE agent_id LIKE 'test-agent%'"))
db_session.execute(text("DELETE FROM agent_registry WHERE id LIKE 'test-agent%'"))
db_session.commit()
except Exception as e:
db_session.rollback()
print(f"E2E cleanup error: {e}")
@pytest.fixture(scope="function")
def mock_llm_streaming():
"""
Mock LLM streaming response for E2E tests.
Returns an async generator that yields streaming chunks.
"""
async def stream_completion(*args, **kwargs):
"""Mock streaming completion with test response."""
chunks = [
"Test ",
"response ",
"chunk 1",
"Test ",
"response ",
"chunk 2",
"Test ",
"response ",
"chunk 3"
]
for chunk in chunks:
yield {
"choices": [{
"delta": {"content": chunk},
"finish_reason": None
}],
"usage": None
}
# Final chunk with finish_reason
yield {
"choices": [{
"delta": {},
"finish_reason": "stop"
}],
"usage": {
"prompt_tokens": 10,
"completion_tokens": 20,
"total_tokens": 30
}
}
return stream_completion
@pytest.fixture(scope="function")
def mock_llm_streaming_error():
"""
Mock LLM streaming error for E2E error path tests.
"""
async def stream_completion_error(*args, **kwargs):
"""Mock streaming completion with error."""
yield {
"choices": [{
"delta": {"content": "Initial chunk"},
"finish_reason": None
}],
"usage": None
}
# Simulate LLM API error
raise Exception("LLM API error: rate limit exceeded")
return stream_completion_error
@pytest.fixture(scope="function")
def mock_websocket():
"""
Mock WebSocket manager for E2E tests.
Mocks WebSocket notifications for agent status updates and execution events.
"""
with patch('core.governance_cache.WebSocketManager') as mock_ws_class:
mock_ws_instance = MagicMock()
mock_ws_instance.notify_agent_status = MagicMock()
mock_ws_instance.notify_execution_start = MagicMock()
mock_ws_instance.notify_execution_complete = MagicMock()
mock_ws_instance.notify_execution_failed = MagicMock()
mock_ws_class.return_value = mock_ws_instance
yield mock_ws_instance
@pytest.fixture(scope="function")
def e2e_client(client, e2e_db_session, mock_websocket):
"""
E2E test client with all necessary mocks.
Combines TestClient with database session, WebSocket mocks,
and authentication bypass for comprehensive E2E testing.
"""
yield client
@pytest.fixture(scope="function")
def execution_id():
"""
Generate unique execution ID for E2E tests.
"""
return str(uuid.uuid4())
# E2E Test Helpers
def assert_episode_created(db_session: Session, agent_id: str, expected_count: int = 1):
"""
Assert that episodes were created for agent execution.
Args:
db_session: Database session
agent_id: Agent ID to check
expected_count: Expected number of episodes (default: 1)
"""
episodes = db_session.query(AgentEpisode).filter(
AgentEpisode.agent_id == agent_id
).all()
assert len(episodes) == expected_count, f"Expected {expected_count} episodes, got {len(episodes)}"
return episodes
def assert_execution_logged(db_session: Session, execution_id: str, expected_status: str = "completed"):
"""
Assert that execution was logged with expected status.
Args:
db_session: Database session
execution_id: Execution ID to check
expected_status: Expected execution status (default: "completed")
"""
execution = db_session.query(AgentExecution).filter(
AgentExecution.id == execution_id
).first()
assert execution is not None, f"Execution {execution_id} not found"
assert execution.status == expected_status, f"Expected status {expected_status}, got {execution.status}"
return execution
def assert_segments_created(db_session: Session, episode_id: str, min_count: int = 1):
"""
Assert that episode segments were created.
Args:
db_session: Database session
episode_id: Episode ID to check
min_count: Minimum number of segments expected (default: 1)
"""
segments = db_session.query(EpisodeSegment).filter(
EpisodeSegment.episode_id == episode_id
).all()
assert len(segments) >= min_count, f"Expected at least {min_count} segments, got {len(segments)}"
return segments
|