trading/tests/services/kevin_signal_bridge/test_aggregator.py

"""Tests for the MentionAggregator — multi-mention window + conviction boost."""

from datetime import datetime, timedelta, timezone
from decimal import Decimal

import pytest
from sqlalchemy.ext.asyncio import AsyncSession

from services.kevin_signal_bridge.aggregator import MentionAggregator
from shared.models.meet_kevin import (
    KevinAnalysis,
    KevinChannel,
    KevinMarketOutlook,
    KevinStockMention,
    KevinTickerAction,
    KevinTimeHorizon,
    KevinVideo,
    KevinVideoStatus,
)


async def _seed_channel_video(
    session: AsyncSession, suffix: str = "1"
) -> tuple[int, int]:
    channel = KevinChannel(youtube_channel_id=f"UC{suffix}", title="t")
    session.add(channel)
    await session.flush()
    video = KevinVideo(
        channel_id=channel.id,
        youtube_video_id=f"v{suffix}",
        title="t",
        published_at=datetime.now(timezone.utc),
        status=KevinVideoStatus.ANALYZED,
    )
    session.add(video)
    await session.flush()
    analysis = KevinAnalysis(
        video_id=video.id,
        model="m",
        prompt_version="v1",
        market_outlook_direction=KevinMarketOutlook.NEUTRAL,
        market_outlook_reasoning="x",
        summary="x",
        prompt_tokens=10,
        completion_tokens=10,
        cost_usd=Decimal("0.01"),
    )
    session.add(analysis)
    await session.flush()
    return video.id, analysis.id


def _factory(session: AsyncSession):
    """Return a session_factory that returns the same session."""

    class _StaticSessionFactory:
        async def __aenter__(self):
            return session

        async def __aexit__(self, *args):
            pass

    def factory():
        return _StaticSessionFactory()

    return factory


async def _insert_mention(
    session: AsyncSession,
    video_id: int,
    analysis_id: int,
    symbol: str,
    conviction: str,
    when: datetime,
    action: KevinTickerAction = KevinTickerAction.BUY,
) -> KevinStockMention:
    m = KevinStockMention(
        video_id=video_id,
        analysis_id=analysis_id,
        symbol=symbol,
        action=action,
        conviction=Decimal(conviction),
        time_horizon=KevinTimeHorizon.WEEKS,
        rationale_quote="x",
    )
    session.add(m)
    await session.flush()
    # Override created_at
    m.created_at = when
    session.add(m)
    await session.flush()
    return m


@pytest.mark.asyncio
async def test_aggregator_returns_all_unseen_mentions(db_session: AsyncSession):
    video_id, analysis_id = await _seed_channel_video(db_session)
    now = datetime.now(timezone.utc)
    m1 = await _insert_mention(
        db_session, video_id, analysis_id, "NVDA", "0.7", now - timedelta(hours=10)
    )
    m2 = await _insert_mention(
        db_session, video_id, analysis_id, "NVDA", "0.7", now - timedelta(hours=5)
    )
    m3 = await _insert_mention(
        db_session, video_id, analysis_id, "INTC", "0.7", now - timedelta(hours=2)
    )

    agg = MentionAggregator(session_factory=_factory(db_session))
    pending = await agg.fetch_pending(since_id=0)
    ids = {p.id for p in pending}
    assert ids == {m1.id, m2.id, m3.id}


@pytest.mark.asyncio
async def test_aggregator_applies_conviction_boost(db_session: AsyncSession):
    video_id, analysis_id = await _seed_channel_video(db_session, "boost")
    now = datetime.now(timezone.utc)
    await _insert_mention(
        db_session, video_id, analysis_id, "NVDA", "0.7", now - timedelta(hours=10)
    )
    m2 = await _insert_mention(
        db_session, video_id, analysis_id, "NVDA", "0.7", now - timedelta(hours=5)
    )

    agg = MentionAggregator(
        session_factory=_factory(db_session),
        window_hours=48,
        boost_per_repeat=Decimal("0.05"),
    )
    pending = await agg.fetch_pending(since_id=0)
    # The second NVDA mention should have effective_conviction = 0.7 + 0.05 = 0.75
    by_id = {p.id: p for p in pending}
    assert by_id[m2.id].effective_conviction == Decimal("0.75")


@pytest.mark.asyncio
async def test_aggregator_caps_boost_at_max(db_session: AsyncSession):
    video_id, analysis_id = await _seed_channel_video(db_session, "cap")
    now = datetime.now(timezone.utc)
    # 6 mentions -> 5 extras -> boost 5*0.05 = 0.25 -> capped at 0.20
    last_m = None
    for i in range(6):
        last_m = await _insert_mention(
            db_session,
            video_id,
            analysis_id,
            "NVDA",
            "0.7",
            now - timedelta(hours=20 - i),
        )

    agg = MentionAggregator(
        session_factory=_factory(db_session),
        window_hours=48,
        boost_per_repeat=Decimal("0.05"),
        max_boost=Decimal("0.20"),
    )
    pending = await agg.fetch_pending(since_id=0)
    by_id = {p.id: p for p in pending}
    assert last_m is not None
    # capped at 0.7 + 0.20 = 0.90
    assert by_id[last_m.id].effective_conviction == Decimal("0.90")


@pytest.mark.asyncio
async def test_aggregator_excludes_already_processed(db_session: AsyncSession):
    video_id, analysis_id = await _seed_channel_video(db_session, "excl")
    now = datetime.now(timezone.utc)
    m1 = await _insert_mention(
        db_session, video_id, analysis_id, "NVDA", "0.7", now - timedelta(hours=10)
    )
    m2 = await _insert_mention(
        db_session, video_id, analysis_id, "INTC", "0.7", now - timedelta(hours=5)
    )

    agg = MentionAggregator(session_factory=_factory(db_session))
    pending = await agg.fetch_pending(since_id=m1.id)
    ids = {p.id for p in pending}
    assert m1.id not in ids
    assert m2.id in ids