topoteretes--cognee/cognee/tests/unit/api/test_api_error_responses.py

"""
Consolidated tests for API endpoint error responses and basic endpoint functionality.

Tests cover:
- Error handling (PipelineRunErrored → 422, PermissionDenied → 403, etc.)
- Basic endpoint functionality (happy path with mocked backends)
"""

import importlib
from types import SimpleNamespace
from uuid import uuid4
from unittest.mock import AsyncMock

import pytest

from fastapi import FastAPI
from fastapi.testclient import TestClient

from cognee.modules.users.methods import get_authenticated_user
from cognee.modules.pipelines.models import PipelineRunErrored, PipelineRunCompleted
from cognee.modules.users.exceptions.exceptions import PermissionDeniedError
from cognee.infrastructure.llm.exceptions import LLMPaymentRequiredError
from cognee.api.v1.add.routers.get_add_router import get_add_router
from cognee.api.v1.cognify.routers.get_cognify_router import get_cognify_router
from cognee.api.v1.memify.routers.get_memify_router import get_memify_router
from cognee.api.v1.remember.routers.get_remember_router import get_remember_router
from cognee.api.v1.search.routers.get_search_router import get_search_router
from cognee.api.v1.update.routers.get_update_router import get_update_router


MOCK_USER = SimpleNamespace(id=uuid4(), email="test@example.com", is_active=True, tenant_id=uuid4())
MOCK_DATASET_ID = uuid4()
MOCK_PIPELINE_RUN_ID = uuid4()


def _make_completed(**kwargs):
    defaults = dict(
        pipeline_run_id=MOCK_PIPELINE_RUN_ID,
        dataset_id=MOCK_DATASET_ID,
        dataset_name="test_dataset",
        status="completed",
    )
    defaults.update(kwargs)
    return PipelineRunCompleted(**defaults)


def _make_errored(error="pipeline failed", **kwargs):
    defaults = dict(
        pipeline_run_id=MOCK_PIPELINE_RUN_ID,
        dataset_id=MOCK_DATASET_ID,
        dataset_name="test_dataset",
        status="errored",
        error=error,
    )
    defaults.update(kwargs)
    return PipelineRunErrored(**defaults)


@pytest.fixture
def app():
    app = FastAPI()
    app.include_router(get_add_router(), prefix="/add")
    app.include_router(get_cognify_router(), prefix="/cognify")
    app.include_router(get_memify_router(), prefix="/memify")
    app.include_router(get_remember_router(), prefix="/remember")
    app.include_router(get_search_router(), prefix="/search")
    app.include_router(get_update_router(), prefix="/update")

    async def override_user():
        return MOCK_USER

    app.dependency_overrides[get_authenticated_user] = override_user
    return app


@pytest.fixture
def client(app):
    return TestClient(app)


# ---------------------------------------------------------------------------
# Add endpoint
# ---------------------------------------------------------------------------


class TestAddEndpoint:
    def test_add_missing_dataset_returns_400(self, client):
        resp = client.post(
            "/add",
            data={},
            files={"data": ("x.txt", b"hello", "text/plain")},
        )
        assert resp.status_code == 400

    def test_add_pipeline_errored_returns_500(self, client):
        import cognee.api.v1.add as add_pkg

        add_pkg.add = AsyncMock(return_value=_make_errored())

        resp = client.post(
            "/add",
            data={"datasetName": "test_dataset"},
            files={"data": ("x.txt", b"hello", "text/plain")},
        )
        assert resp.status_code == 500
        body = resp.json()
        assert body["error"] == "Pipeline run errored"
        assert "detail" in body

    def test_add_success_returns_200(self, client):
        import cognee.api.v1.add as add_pkg

        completed = _make_completed()
        add_pkg.add = AsyncMock(return_value=completed)

        resp = client.post(
            "/add",
            data={"datasetName": "test_dataset"},
            files={"data": ("test.txt", b"Cognee is an AI memory platform.", "text/plain")},
        )
        assert resp.status_code == 200
        body = resp.json()
        assert body["status"] == "completed"
        assert body["dataset_name"] == "test_dataset"

    def test_add_internal_error_returns_500(self, client):
        import cognee.api.v1.add as add_pkg

        add_pkg.add = AsyncMock(side_effect=RuntimeError("unexpected"))

        resp = client.post(
            "/add",
            data={"datasetName": "test_dataset"},
            files={"data": ("x.txt", b"hello", "text/plain")},
        )
        assert resp.status_code == 500
        assert resp.json()["error"] == "Internal server error"


# ---------------------------------------------------------------------------
# Remember endpoint
# ---------------------------------------------------------------------------


class TestRememberEndpoint:
    def test_remember_passes_chunk_size(self, client, monkeypatch):
        remember_pkg = importlib.import_module("cognee.api.v1.remember")

        class RememberResultStub:
            status = "completed"

            def to_dict(self):
                return {"status": "completed", "dataset_name": "test_dataset"}

        remember = AsyncMock(return_value=RememberResultStub())
        monkeypatch.setattr(remember_pkg, "remember", remember)

        resp = client.post(
            "/remember",
            data={"datasetName": "test_dataset", "chunk_size": "42"},
            files={"data": ("test.txt", b"Cognee is an AI memory platform.", "text/plain")},
        )

        assert resp.status_code == 200
        remember.assert_awaited_once()
        assert remember.await_args.kwargs["chunk_size"] == 42


# ---------------------------------------------------------------------------
# Cognify endpoint
# ---------------------------------------------------------------------------


class TestCognifyEndpoint:
    def test_cognify_missing_datasets_returns_400(self, client):
        resp = client.post("/cognify", json={})
        assert resp.status_code == 400

    def test_cognify_pipeline_errored_returns_500(self, client):
        import cognee.api.v1.cognify as cognify_pkg

        cognify_pkg.cognify = AsyncMock(return_value={"run": _make_errored()})

        resp = client.post(
            "/cognify",
            json={"datasets": ["test_dataset"], "run_in_background": False},
        )
        assert resp.status_code == 500
        body = resp.json()
        assert body["error"] == "Pipeline run errored"

    def test_cognify_success_returns_200(self, client):
        import cognee.api.v1.cognify as cognify_pkg

        completed = _make_completed()
        cognify_pkg.cognify = AsyncMock(return_value={str(MOCK_DATASET_ID): completed})

        resp = client.post(
            "/cognify",
            json={"datasets": ["test_dataset"], "run_in_background": False},
        )
        assert resp.status_code == 200

    def test_cognify_passes_chunk_size(self, client, monkeypatch):
        import cognee.api.v1.cognify as cognify_pkg

        completed = _make_completed()
        cognify = AsyncMock(return_value={str(MOCK_DATASET_ID): completed})
        monkeypatch.setattr(cognify_pkg, "cognify", cognify)

        resp = client.post(
            "/cognify",
            json={"datasets": ["test_dataset"], "chunk_size": 42},
        )

        assert resp.status_code == 200
        cognify.assert_awaited_once()
        assert cognify.await_args.kwargs["chunk_size"] == 42

    def test_cognify_internal_error_returns_500(self, client):
        import cognee.api.v1.cognify as cognify_pkg

        cognify_pkg.cognify = AsyncMock(side_effect=RuntimeError("unexpected"))

        resp = client.post(
            "/cognify",
            json={"datasets": ["test_dataset"]},
        )
        assert resp.status_code == 500
        assert resp.json()["error"] == "Internal server error"

    def test_cognify_llm_payment_required_returns_402(self, client):
        import cognee.api.v1.cognify as cognify_pkg

        cognify_pkg.cognify = AsyncMock(side_effect=LLMPaymentRequiredError())

        resp = client.post(
            "/cognify",
            json={"datasets": ["test_dataset"]},
        )
        assert resp.status_code == 402
        assert resp.json()["error"] == "Token budget exhausted"


# ---------------------------------------------------------------------------
# Search endpoint
# ---------------------------------------------------------------------------


class TestSearchEndpoint:
    def test_search_permission_denied_returns_403(self, client):
        import cognee.api.v1.search as search_pkg

        search_pkg.search = AsyncMock(side_effect=PermissionDeniedError("no access to dataset"))

        resp = client.post(
            "/search",
            json={
                "search_type": "GRAPH_COMPLETION",
                "query": "What is Cognee?",
            },
        )
        assert resp.status_code == 403
        body = resp.json()
        assert body["error"] == "Permission denied"

    def test_search_success_returns_200(self, client):
        import cognee.api.v1.search as search_pkg

        search_pkg.search = AsyncMock(
            return_value=[
                {
                    "search_result": "Cognee is an AI memory platform.",
                    "dataset_id": str(MOCK_DATASET_ID),
                    "dataset_name": "test_dataset",
                }
            ]
        )

        resp = client.post(
            "/search",
            json={
                "search_type": "GRAPH_COMPLETION",
                "query": "What is Cognee?",
            },
        )
        assert resp.status_code == 200
        body = resp.json()
        assert len(body) == 1
        assert "Cognee" in body[0]["search_result"]

    def test_search_internal_error_returns_500(self, client):
        import cognee.api.v1.search as search_pkg

        search_pkg.search = AsyncMock(side_effect=RuntimeError("unexpected"))

        resp = client.post(
            "/search",
            json={"query": "test"},
        )
        assert resp.status_code == 500
        assert resp.json()["error"] == "Internal server error"

    def test_search_llm_payment_required_returns_402(self, client):
        import cognee.api.v1.search as search_pkg

        search_pkg.search = AsyncMock(side_effect=LLMPaymentRequiredError())

        resp = client.post(
            "/search",
            json={"search_type": "GRAPH_COMPLETION", "query": "test"},
        )
        assert resp.status_code == 402
        assert resp.json()["error"] == "Token budget exhausted"


# ---------------------------------------------------------------------------
# Memify endpoint
# ---------------------------------------------------------------------------


class TestMemifyEndpoint:
    def test_memify_missing_dataset_returns_400(self, client):
        resp = client.post("/memify", json={})
        assert resp.status_code == 400

    def test_memify_pipeline_errored_returns_500(self, client):
        import cognee.modules.memify as memify_pkg

        memify_pkg.memify = AsyncMock(return_value=_make_errored(error="memify failed"))

        resp = client.post(
            "/memify",
            json={"dataset_name": "test_dataset", "run_in_background": False},
        )
        assert resp.status_code == 500
        body = resp.json()
        assert body["error"] == "Pipeline run errored"

    def test_memify_success_returns_200(self, client):
        import cognee.modules.memify as memify_pkg

        completed = _make_completed()
        memify_pkg.memify = AsyncMock(return_value=completed.model_dump())

        resp = client.post(
            "/memify",
            json={"dataset_name": "test_dataset", "run_in_background": False},
        )
        assert resp.status_code == 200
        body = resp.json()
        assert body["status"] == "completed"

    def test_memify_internal_error_returns_500(self, client):
        import cognee.modules.memify as memify_pkg

        memify_pkg.memify = AsyncMock(side_effect=RuntimeError("unexpected"))

        resp = client.post(
            "/memify",
            json={"dataset_name": "test_dataset"},
        )
        assert resp.status_code == 500
        assert resp.json()["error"] == "Internal server error"


# ---------------------------------------------------------------------------
# Update endpoint
# ---------------------------------------------------------------------------


class TestUpdateEndpoint:
    def test_update_pipeline_errored_returns_500(self, client):
        import cognee.api.v1.update as update_pkg

        update_pkg.update = AsyncMock(return_value={"run": _make_errored(error="update failed")})

        resp = client.patch(
            "/update",
            params={"data_id": str(uuid4()), "dataset_id": str(uuid4())},
            files={"data": ("x.txt", b"hello", "text/plain")},
            data={"node_set": ""},
        )
        assert resp.status_code == 500
        body = resp.json()
        assert body["error"] == "Pipeline run errored"

    def test_update_success_returns_200(self, client):
        import cognee.api.v1.update as update_pkg

        completed = _make_completed()
        update_pkg.update = AsyncMock(return_value={str(MOCK_DATASET_ID): completed})

        resp = client.patch(
            "/update",
            params={"data_id": str(uuid4()), "dataset_id": str(uuid4())},
            files={"data": ("updated.txt", b"updated content", "text/plain")},
            data={"node_set": ""},
        )
        assert resp.status_code == 200

    def test_update_internal_error_returns_500(self, client):
        import cognee.api.v1.update as update_pkg

        update_pkg.update = AsyncMock(side_effect=RuntimeError("unexpected"))

        resp = client.patch(
            "/update",
            params={"data_id": str(uuid4()), "dataset_id": str(uuid4())},
            files={"data": ("x.txt", b"hello", "text/plain")},
            data={"node_set": ""},
        )
        assert resp.status_code == 500
        assert resp.json()["error"] == "Internal server error"