presenton/servers/fastapi/tests/test_stream_presentation.py
2025-05-10 21:27:30 +05:45

88 lines
3 KiB
Python

from unittest.mock import patch, MagicMock
from fastapi.testclient import TestClient
import json
from api.main import app
from api.models import SessionModel
client = TestClient(app)
def mock_stream_response():
return [
{"type": "start", "data": "Starting presentation generation"},
{"type": "slide_start", "data": "Generating slide 1"},
{"type": "slide_end", "data": "Completed slide 1"},
{"type": "slide_start", "data": "Generating slide 2"},
{"type": "slide_end", "data": "Completed slide 2"},
{"type": "end", "data": "Presentation generation complete"}
]
@patch('api.routers.presentation.handlers.generate_stream.PresentationGenerateStreamHandler.get')
def test_presentation_flow(mock_stream):
# Setup mock for streaming response
mock_stream.return_value = mock_stream_response()
# Step 1: Upload document
test_file_path = "tests/assets/impact_of_llms.pdf"
with open(test_file_path, "rb") as f:
upload_response = client.post(
"/ppt/files/upload",
files={"documents": ("impact_of_llms.pdf", f, "application/pdf")}
)
assert upload_response.status_code == 200
upload_json = upload_response.json()
assert "documents" in upload_json
assert len(upload_json["documents"]) > 0
# Step 2: Decompose uploaded document
decompose_response = client.post(
"/ppt/files/decompose",
json={"documents": upload_json["documents"]}
)
assert decompose_response.status_code == 200
decompose_json = decompose_response.json()
assert "documents" in decompose_json
assert upload_json["documents"][0] in decompose_json["documents"]
# Step 3: Create presentation
create_response = client.post(
"/ppt/create",
json={
"prompt": "Create a presentation about LLMs",
"n_slides": 2,
"language": "en",
"documents": upload_json["documents"]
}
)
assert create_response.status_code == 200
presentation = create_response.json()
presentation_id = presentation["id"]
# Step 4: Generate presentation data and get session
gen_data_response = client.post(
"/ppt/generate/data",
json={
"presentation_id": presentation_id,
"theme": None,
"images": None,
"watermark": True,
"titles": ["Introduction to LLMs", "Impact of LLMs"]
}
)
assert gen_data_response.status_code == 200
session_data = gen_data_response.json()
assert "session" in session_data
# Step 5: Stream presentation generation
stream_response = client.get(f"/ppt/generate/stream?presentation_id={presentation_id}&session={session_data['session']}")
assert stream_response.status_code == 200
# Verify stream events
events = list(mock_stream_response())
assert len(events) > 0
assert events[0]["type"] == "start"
assert events[-1]["type"] == "end"
assert any(event["type"] == "slide_start" for event in events)
assert any(event["type"] == "slide_end" for event in events)