backend/tests/test_e2e_zus.py

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(
        "rag_pipeline.vector_store.get_vector_store"
    ) as mock_pinecone:
        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 dla ZUS. " * 100
            + "\n\n| Element | Koszt |\n|---|---|\n| Wentylacja | 5000 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_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"]