discovery-by-llm/backend/tests/test_orchestrator.py

"""End-to-end pipeline tests using deterministic service fakes."""

import asyncio
from collections.abc import AsyncGenerator

from app.config import Settings
from app.domain.models import (
    ConversationTurn,
    Familiarity,
    Intent,
    PreviousRecommendation,
    RerankSelection,
    Track,
    TrackCandidate,
)
from app.pipeline.event import PipelineEvent, PipelineTrackEvent
from app.pipeline.orchestrator import RecommendationPipeline
from app.ports.protocols import RecommenderOutputError, TimeRange


class FakeCatalog:
    """Return exact tracks and a configurable known-track sample."""

    def __init__(self, tracks: tuple[Track, ...], known_tracks: tuple[Track, ...] = ()) -> None:
        """Index tracks by title and expose taste-call counters."""
        self.tracks_by_title = {track.title: track for track in tracks}
        self.known_tracks = known_tracks
        self.search_call_count = 0
        self.taste_call_count = 0

    async def search_tracks(self, query: str, limit: int = 10) -> list[Track]:
        """Resolve an exact fielded title and miss bare fallbacks."""
        self.search_call_count += 1
        if 'track:"' not in query:
            return []
        title = query.split('track:"', 1)[1].split('"', 1)[0]
        track = self.tracks_by_title.get(title)
        return [track] if track is not None else []

    async def fetch_top_artists(self, time_range: TimeRange, limit: int) -> list[str]:
        """Return one stable taste artist."""
        self.taste_call_count += 1
        return ["Taste Artist"]

    async def fetch_top_tracks(self, time_range: TimeRange, limit: int) -> list[Track]:
        """Return no top tracks."""
        self.taste_call_count += 1
        return []

    async def fetch_saved_tracks(self, limit: int) -> list[Track]:
        """Return the configured known tracks."""
        self.taste_call_count += 1
        return list(self.known_tracks)


class FakeRecommender:
    """Return fixed intents and either selections or structured failures."""

    def __init__(
        self,
        intents: list[Intent],
        selection_ids: tuple[str, ...] = (),
        failure_count: int = 0,
    ) -> None:
        """Store deterministic outputs for successive calls."""
        self.intents = intents
        self.selection_ids = selection_ids
        self.failure_count = failure_count
        self.rerank_call_count = 0

    async def create_intent(
        self,
        query: str,
        history: tuple[ConversationTurn, ...],
        previous_recommendations: tuple[PreviousRecommendation, ...],
        taste_summary: str,
        candidate_count: int,
    ) -> Intent:
        """Return the next fixed intent."""
        return self.intents.pop(0)

    async def stream_rerank(
        self,
        intent: Intent,
        grounded_tracks: tuple[Track, ...],
        taste_summary: str,
        history: tuple[ConversationTurn, ...],
        selection_count: int,
        correction: str | None = None,
    ) -> AsyncGenerator[RerankSelection]:
        """Stream configured ids or fail before yielding."""
        self.rerank_call_count += 1
        if self.failure_count:
            self.failure_count -= 1
            raise RecommenderOutputError("invalid test output")
        selected_ids = self.selection_ids or tuple(track.id for track in grounded_tracks)
        for track_id in selected_ids[:selection_count]:
            yield RerankSelection(track_id, f"Reason for {track_id}")


def test_event_order_and_rerank_ids_stay_inside_grounded_pool() -> None:
    async def run() -> None:
        first = _track("first", "First Song")
        second = _track("second", "Second Song")
        catalog = FakeCatalog((first, second))
        recommender = FakeRecommender([_intent(first, second)], ("second", "first"))

        events = await _run_pipeline(catalog, recommender)

        assert [event.type for event in events] == ["metadata", "track", "track", "done"]
        track_events = [event for event in events if isinstance(event, PipelineTrackEvent)]
        assert [event.track.id for event in track_events] == ["second", "first"]
        assert {event.track.id for event in track_events} <= {"first", "second"}

    asyncio.run(run())


def test_rerank_fallback_warns_then_streams_grounded_order() -> None:
    async def run() -> None:
        first = _track("first", "First Song")
        second = _track("second", "Second Song")
        catalog = FakeCatalog((first, second))
        recommender = FakeRecommender([_intent(first, second)], failure_count=2)

        events = await _run_pipeline(catalog, recommender)

        assert [event.type for event in events] == [
            "metadata",
            "warning",
            "track",
            "track",
            "done",
        ]
        assert recommender.rerank_call_count == 2
        track_events = [event for event in events if isinstance(event, PipelineTrackEvent)]
        assert [event.track.id for event in track_events] == ["first", "second"]

    asyncio.run(run())


def test_new_familiarity_excludes_known_track_ids() -> None:
    async def run() -> None:
        known = _track("known", "Known Song")
        new = _track("new", "New Song")
        catalog = FakeCatalog((known, new), known_tracks=(known,))
        intent = _intent(known, new, familiarity=Familiarity.NEW)
        recommender = FakeRecommender([intent])

        events = await _run_pipeline(catalog, recommender, rerank_count=1)

        track_events = [event for event in events if isinstance(event, PipelineTrackEvent)]
        assert [event.track.id for event in track_events] == ["new"]

    asyncio.run(run())


def test_refinement_reuses_last_grounded_pool_and_cached_taste() -> None:
    async def run() -> None:
        track = _track("first", "First Song")
        initial = _intent(track)
        refinement = _intent(track, is_refinement=True)
        catalog = FakeCatalog((track,))
        recommender = FakeRecommender([initial, refinement])
        pipeline = _pipeline(recommender, rerank_count=1)

        await _collect(pipeline, catalog, "first request")
        initial_search_calls = catalog.search_call_count
        initial_taste_calls = catalog.taste_call_count
        await _collect(pipeline, catalog, "refine it")

        assert catalog.search_call_count == initial_search_calls
        assert catalog.taste_call_count == initial_taste_calls

    asyncio.run(run())


async def _run_pipeline(
    catalog: FakeCatalog,
    recommender: FakeRecommender,
    rerank_count: int = 2,
) -> list[PipelineEvent]:
    return await _collect(_pipeline(recommender, rerank_count=rerank_count), catalog, "query")


def _pipeline(recommender: FakeRecommender, rerank_count: int) -> RecommendationPipeline:
    return RecommendationPipeline(
        recommender,
        Settings(
            rerank_count=rerank_count,
            rerank_pool_buffer=0,
            grounding_floor=1,
            grounding_concurrency=2,
            request_deadline_seconds=1.0,
        ),
    )


async def _collect(
    pipeline: RecommendationPipeline,
    catalog: FakeCatalog,
    query: str,
) -> list[PipelineEvent]:
    return [
        event
        async for event in pipeline.stream(
            "session",
            "request",
            catalog,
            query,
            (),
            (),
        )
    ]


def _intent(
    *tracks: Track,
    familiarity: Familiarity = Familiarity.MIX,
    is_refinement: bool = False,
) -> Intent:
    return Intent(
        mood=("focused",),
        activity=None,
        era=(),
        languages=(),
        genres=("electronic",),
        familiarity=familiarity,
        is_refinement=is_refinement,
        intent_summary="Focused electronic discovery.",
        candidates=tuple(
            TrackCandidate(title=track.title, artist=track.artists[0]) for track in tracks
        ),
    )


def _track(track_id: str, title: str) -> Track:
    return Track(
        id=track_id,
        uri=f"spotify:track:{track_id}",
        title=title,
        artists=("Artist",),
        album_name="Album",
        album_art_url=None,
        external_url=None,
    )


def test_below_floor_pool_streams_partial_results_after_warning() -> None:
    async def run() -> None:
        found = _track("found", "Found Song")
        missing = _track("missing", "Missing Song")
        catalog = FakeCatalog((found,))
        recommender = FakeRecommender([_intent(found, missing)])
        pipeline = RecommendationPipeline(
            recommender,
            Settings(
                rerank_count=2,
                rerank_pool_buffer=0,
                grounding_floor=2,
                grounding_concurrency=2,
                request_deadline_seconds=1.0,
            ),
        )

        events = await _collect(pipeline, catalog, "query")

        assert [event.type for event in events] == ["metadata", "warning", "track", "done"]
        warning = events[1]
        assert warning.type == "warning"
        assert warning.code == "partial_results"

    asyncio.run(run())


def test_empty_pool_is_a_terminal_error() -> None:
    async def run() -> None:
        missing = _track("missing", "Missing Song")
        catalog = FakeCatalog(())
        recommender = FakeRecommender([_intent(missing)])

        events = await _run_pipeline(catalog, recommender)

        assert [event.type for event in events] == ["metadata", "error"]
        error = events[-1]
        assert error.type == "error"
        assert error.code == "no_grounded_results"

    asyncio.run(run())