tests/test_generate_focus_selection.py

from app.pipeline.nodes.generate import _select_chunks_for_prompt


def _chunk(title: str, section: str = "Overview") -> dict:
    return {
        "text": f"Info about {title}",
        "metadata": {
            "source_title": title,
            "section": section,
            "source_url": "",
            "source_type": "project",
        },
    }


def test_select_chunks_prefers_explicitly_mentioned_source_title() -> None:
    chunks = [
        _chunk("Sorting Demo"),
        _chunk("EchoEcho"),
        _chunk("Donut.asm"),
    ]

    selected = _select_chunks_for_prompt(
        "Is the source code on GitHub up-to-date with the Sorting Demo live demo?",
        chunks,
    )

    assert selected
    assert all(c["metadata"]["source_title"] == "Sorting Demo" for c in selected)


def test_select_chunks_falls_back_to_top_ranked_when_no_title_match() -> None:
    chunks = [_chunk("A"), _chunk("B"), _chunk("C")]

    selected = _select_chunks_for_prompt("What technologies are used?", chunks)

    assert selected == chunks