Spaces:
Running
Running
File size: 3,897 Bytes
2754b82 | 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 | import pytest
from httpx import AsyncClient
from unittest.mock import patch, MagicMock
@pytest.mark.asyncio
async def test_e2e_zus_lifecycle(async_client: AsyncClient, auth_headers):
# Mock LLM calls
with patch("endpoints.projects.get_llm") as mock_compile_llm, patch(
"agents.helpers.wizard_node"
) as mock_wizard_node, patch(
"agents.helpers.generate_section_light"
) as mock_generate_section_light, patch(
"rag_pipeline.vector_store.get_vector_store"
) as mock_pinecone:
mock_generate_section_light.return_value = "Przykładowa testowa treść"
mock_pinecone.return_value = MagicMock()
mock_pinecone.return_value.as_retriever.return_value.invoke.return_value = [
MagicMock(page_content="Zmockowany kontekst RAG dla ZUS")
]
# Mocking wizard_node return for section generation
valid_mock_content = (
"Przykładowa testowa treść wniosku o dofinansowanie. " * 100
+ "\n\n| Tabela | Koszt |\n|---|---|\n| Zadanie 1 | 1000 PLN |\n\n"
+ "Dodatkowy tekst wypełniający, aby osiągnąć odpowiednią długość. " * 100
)
if len(valid_mock_content) < 5000:
valid_mock_content += "A" * (5000 - len(valid_mock_content))
mock_wizard_node.return_value = {
"messages": [MagicMock(content=valid_mock_content)]
}
mock_generate_section_light.return_value = valid_mock_content
mock_response = MagicMock()
mock_response.content = valid_mock_content
mock_compile_llm.return_value = MagicMock()
mock_compile_llm.return_value.invoke.return_value = mock_response
mock_compile_llm.return_value.return_value = mock_response
# 1. Create a ZUS project
payload = {
"title": "Test ZUS BHP",
"program_name": "Dofinansowanie BHP ZUS",
"program_type": "zus",
"description": "Zakup wentylacji",
}
response = await async_client.post(
"/api/projects", json=payload, headers=auth_headers
)
assert response.status_code == 200
project_id = response.json()["id"]
# 2. Get generated structure (sections) -> there's a fallback to SMART if template not found, or it might be default sections.
response = await async_client.get(
f"/api/projects/{project_id}/sections", headers=auth_headers
)
assert response.status_code == 200
sections = response.json()
assert len(sections) > 0, "Brak sekcji startowych dla ZUS"
# 3. Generate a specific section
payload_gen = {
"section_type": "project_summary",
"prompt_context": "Poprawa bezpieczeństwa bhp",
}
response = await async_client.post(
f"/api/projects/{project_id}/generate-section",
json=payload_gen,
headers=auth_headers,
)
print("GENERATE SECTION RESPONSE:", response.status_code, response.json())
assert response.status_code == 200
assert "Przykładowa testowa treść" in response.json()["content"]
# 4. Generate final document
response_check = await async_client.get(
f"/api/projects/{project_id}/sections", headers=auth_headers
)
print("SECTIONS BEFORE COMPILE FINAL:", response_check.json())
payload_compile = {"approved_only": False}
response = await async_client.post(
f"/api/projects/{project_id}/compile-final",
json=payload_compile,
headers=auth_headers,
)
print("COMPILE FINAL RESPONSE:", response.status_code, response.json())
assert response.status_code == 200
assert "final_markdown" in response.json()
assert "Przykładowa testowa treść" in response.json()["final_markdown"]
|