tests/test_integration.py

"""Integration test — full pipeline flow with mocked stages."""

from pathlib import Path
from unittest.mock import AsyncMock, patch

import pytest

from densefeed.config import DenseFeedConfig, PipelineConfig
from densefeed.models import Item
from densefeed.pipeline import Pipeline, PipelineRegister


def _make_items(n: int = 5) -> list[Item]:
    return [
        Item(
            source="test",
            title=f"Story {i}",
            url=f"https://example.com/{i}",
            summary=f"Summary {i}",
            published_ts=1717400000.0 + i,
        )
        for i in range(n)
    ]


@pytest.mark.asyncio
async def test_pipeline_register_roundtrip(tmp_path: Path):
    """Verify register writes and reads stage artifacts correctly."""
    reg = PipelineRegister("2026-06-04", tmp_path)
    data = [{"source": "test", "title": "hello"}]
    reg.write_jsonl(1, "stories.jsonl", data)

    assert reg.is_complete(1)
    assert not reg.is_complete(2)
    assert reg.find_last_complete() == 1

    loaded = reg.read_jsonl(1, "stories.jsonl")
    assert loaded[0]["title"] == "hello"


@pytest.mark.asyncio
async def test_pipeline_stops_on_failure(tmp_path: Path):
    """Pipeline should stop running stages after a failure."""
    cfg = DenseFeedConfig(pipeline=PipelineConfig(run_dir=str(tmp_path)))
    pipeline = Pipeline(cfg)

    # Mock fetch to succeed, rank to fail
    with (
        patch(
            "densefeed.pipeline.Pipeline._stage_fetch", new_callable=AsyncMock
        ) as mock_fetch,
        patch(
            "densefeed.pipeline.Pipeline._stage_rank", new_callable=AsyncMock
        ) as mock_rank,
    ):
        from densefeed.models import StageResult

        mock_fetch.return_value = StageResult(
            stage=1, name="fetch", success=True, items_out=5
        )
        mock_rank.return_value = StageResult(
            stage=2, name="rank", success=False, error="LLM timeout"
        )

        result = await pipeline.run("2026-06-04")

    # Only fetch should have run (successfully), then rank failed
    assert len(result.stages) == 2
    assert result.stages[0].success
    assert not result.stages[1].success


@pytest.mark.asyncio
async def test_pipeline_single_stage(tmp_path: Path):
    """Pipeline with single_stage=True should stop after first stage."""
    cfg = DenseFeedConfig(pipeline=PipelineConfig(run_dir=str(tmp_path)))
    pipeline = Pipeline(cfg)

    with patch(
        "densefeed.pipeline.Pipeline._stage_fetch", new_callable=AsyncMock
    ) as mock_fetch:
        from densefeed.models import StageResult

        mock_fetch.return_value = StageResult(
            stage=1, name="fetch", success=True, items_out=5
        )

        result = await pipeline.run("2026-06-04", start_from=1, single_stage=True)

    assert len(result.stages) == 1
    assert result.stages[0].name == "fetch"


@pytest.mark.asyncio
async def test_pipeline_resume_from_stage(tmp_path: Path):
    """Starting from stage 3 should skip stages 1 and 2."""
    cfg = DenseFeedConfig(pipeline=PipelineConfig(run_dir=str(tmp_path)))
    pipeline = Pipeline(cfg)

    with patch(
        "densefeed.pipeline.Pipeline._stage_extract", new_callable=AsyncMock
    ) as mock_extract:
        from densefeed.models import StageResult

        mock_extract.return_value = StageResult(
            stage=3, name="extract", success=True, items_out=3
        )

        result = await pipeline.run("2026-06-04", start_from=3, single_stage=True)

    assert len(result.stages) == 1
    assert result.stages[0].name == "extract"
    mock_extract.assert_called_once()