Spaces:
Sleeping
Sleeping
File size: 4,118 Bytes
07ed4f9 | 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 | import pytest
from unittest.mock import AsyncMock, patch
from fastapi.testclient import TestClient
from app.main import app
client = TestClient(app)
@pytest.fixture
def mock_trigger():
with patch("app.routes.ai.trigger_simulation", new_callable=AsyncMock) as mock:
yield mock
@pytest.mark.asyncio
async def test_simulate_ai_endpoint(mock_trigger):
# Mock the response from the AI Agents API
mock_trigger.return_value = {
"status": "success",
"seller_id": "TEST_SELLER",
"executive_plan": {
"summary": "This is a mock plan",
"actions": []
}
}
# We also need to mock `embedding_service.store_insight` since it connects to the DB
with patch("app.routes.ai.embedding_service.store_insight", new_callable=AsyncMock) as mock_store:
response = client.post(
"/ai/simulate",
headers={"Authorization": "Bearer dev-api-key"},
json={
"seller_id": "TEST_SELLER",
"time_window_start": "2026-02-01",
"time_window_end": "2026-02-15",
"snapshot_data": {"test": "data"}
}
)
assert response.status_code == 200
data = response.json()
assert data["status"] == "success"
assert "executive_plan" in data
# Verify the mock was called correctly
mock_trigger.assert_called_once_with(
seller_id="TEST_SELLER",
time_window_start="2026-02-01",
time_window_end="2026-02-15",
snapshot_data={"test": "data"}
)
# Verify it attempted to save the insight
mock_store.assert_called_once()
@pytest.fixture
def mock_stream_trigger():
with patch("app.routes.ai.trigger_simulation_stream") as mock:
yield mock
@pytest.mark.asyncio
async def test_simulate_ai_stream_endpoint(mock_stream_trigger):
# Mock an async generator
async def mock_generator():
yield b'data: {"content": "Hello"}\n\n'
yield b'data: {"content": " World"}\n\n'
yield b'data: {"status": "done"}\n\n'
mock_stream_trigger.return_value = mock_generator()
with client.stream("POST", "/ai/simulate/stream",
headers={"Authorization": "Bearer dev-api-key"},
json={
"seller_id": "TEST_SELLER",
"time_window_start": "2026-02-01",
"time_window_end": "2026-02-15",
"snapshot_data": {}
}) as response:
assert response.status_code == 200
assert response.headers["content-type"] == "text/event-stream; charset=utf-8"
chunks = list(response.iter_bytes())
assert len(chunks) == 3
assert b'Hello' in chunks[0]
assert b'World' in chunks[1]
assert b'done' in chunks[2]
@pytest.fixture
def mock_whatif_stream_trigger():
with patch("app.routes.ai.trigger_whatif_stream") as mock:
yield mock
@pytest.mark.asyncio
async def test_simulate_ai_whatif_stream_endpoint(mock_whatif_stream_trigger):
# Mock an async generator
async def mock_generator():
yield b'data: {"content": "Simulation"}\n\n'
yield b'data: {"content": " Results"}\n\n'
yield b'data: {"status": "done"}\n\n'
mock_whatif_stream_trigger.return_value = mock_generator()
with client.stream("POST", "/ai/whatif",
headers={"Authorization": "Bearer dev-api-key"},
json={
"seller_id": "TEST_SELLER",
"scenario": "What if I drop my price 10%?"
}) as response:
assert response.status_code == 200
assert response.headers["content-type"] == "text/event-stream; charset=utf-8"
# Read the streamed chunks
chunks = list(response.iter_bytes())
assert len(chunks) == 3
assert b'Simulation' in chunks[0]
assert b'Results' in chunks[1]
assert b'done' in chunks[2]
|