signal-platform/tests/unit/test_ticker_universe_service.py

"""Unit tests for ticker_universe_service bootstrap logic."""

from __future__ import annotations

import json
from collections.abc import AsyncGenerator

import pytest
from sqlalchemy import select
from sqlalchemy.ext.asyncio import AsyncSession, async_sessionmaker, create_async_engine

from app.database import Base
from app.models.settings import SystemSetting
from app.models.ticker import Ticker
from app.services import ticker_service, ticker_universe_service
from app.services.ticker_universe_service import _extract_wiki_symbols, _normalise_symbols

_engine = create_async_engine("sqlite+aiosqlite://", echo=False)
_session_factory = async_sessionmaker(_engine, class_=AsyncSession, expire_on_commit=False)


@pytest.fixture(autouse=True)
async def _setup_tables() -> AsyncGenerator[None, None]:
    async with _engine.begin() as conn:
        await conn.run_sync(Base.metadata.create_all)
    yield
    async with _engine.begin() as conn:
        await conn.run_sync(Base.metadata.drop_all)


@pytest.fixture
async def session() -> AsyncGenerator[AsyncSession, None]:
    async with _session_factory() as s:
        yield s


@pytest.mark.asyncio
async def test_bootstrap_universe_adds_missing_symbols(session: AsyncSession, monkeypatch: pytest.MonkeyPatch):
    session.add(Ticker(symbol="AAPL"))
    await session.commit()

    async def _fake_fetch(_db: AsyncSession, _universe: str) -> tuple[list[str], str]:
        return ["AAPL", "MSFT", "NVDA"], "test"

    monkeypatch.setattr(ticker_universe_service, "fetch_universe_symbols", _fake_fetch)

    result = await ticker_universe_service.bootstrap_universe(session, "sp500")

    assert result["added"] == 2
    assert result["already_tracked"] == 1
    assert result["deleted"] == 0
    assert result["source"] == "test"
    assert set(result["added_symbols"]) == {"MSFT", "NVDA"}

    rows = await session.execute(select(Ticker.symbol).order_by(Ticker.symbol.asc()))
    assert list(rows.scalars().all()) == ["AAPL", "MSFT", "NVDA"]


@pytest.mark.asyncio
async def test_bootstrap_universe_prunes_missing_symbols(session: AsyncSession, monkeypatch: pytest.MonkeyPatch):
    # Index-owned: a row defaults to manual, and manual rows are never pruned.
    session.add_all([
        Ticker(symbol=s, tracking_source=ticker_service.SOURCE_INDEX)
        for s in ("AAPL", "MSFT", "TSLA")
    ])
    await session.commit()

    async def _fake_fetch(_db: AsyncSession, _universe: str) -> tuple[list[str], str]:
        # A live source label: a fallback result must not drive pruning, so a
        # made-up label would make this test vacuous.
        return ["AAPL", "MSFT"], "wikipedia_sp500"

    monkeypatch.setattr(ticker_universe_service, "fetch_universe_symbols", _fake_fetch)

    result = await ticker_universe_service.bootstrap_universe(
        session,
        "sp500",
        prune_missing=True,
    )

    assert result["added"] == 0
    assert result["already_tracked"] == 2
    assert result["untracked"] == 1
    assert result["untracked_symbols"] == ["TSLA"]
    # Nothing is destroyed any more, and the key says so rather than vanishing.
    assert result["deleted"] == 0

    # The row survives the prune; only its scanner membership changed.
    rows = await session.execute(select(Ticker.symbol).order_by(Ticker.symbol.asc()))
    assert list(rows.scalars().all()) == ["AAPL", "MSFT", "TSLA"]
    live = await session.execute(
        ticker_service.active_only(select(Ticker.symbol).order_by(Ticker.symbol.asc()))
    )
    assert list(live.scalars().all()) == ["AAPL", "MSFT"]


@pytest.mark.asyncio
async def test_fetch_universe_symbols_uses_cached_snapshot_when_live_sources_fail(
    session: AsyncSession,
    monkeypatch: pytest.MonkeyPatch,
):
    session.add(
        SystemSetting(
            key="ticker_universe_cache_sp500",
            value=json.dumps({"symbols": ["AAPL", "MSFT"], "source": "test"}),
        )
    )
    await session.commit()

    async def _fake_public(_universe: str):
        return [], ["public failed"], None

    monkeypatch.setattr(ticker_universe_service, "_fetch_universe_symbols_from_public", _fake_public)

    symbols, source = await ticker_universe_service.fetch_universe_symbols(session, "sp500")
    assert symbols == ["AAPL", "MSFT"]
    assert source == "cache"


@pytest.mark.asyncio
async def test_fetch_universe_symbols_uses_seed_when_live_and_cache_fail(
    session: AsyncSession,
    monkeypatch: pytest.MonkeyPatch,
):
    async def _fake_public(_universe: str):
        return [], ["public failed"], None

    monkeypatch.setattr(ticker_universe_service, "_fetch_universe_symbols_from_public", _fake_public)

    symbols, source = await ticker_universe_service.fetch_universe_symbols(session, "sp500")
    assert "AAPL" in symbols
    assert len(symbols) > 10
    assert source == "seed"


# Snippet shaped like current Wikipedia NyseSymbol / exchange link markup.
_SAMPLE_WIKI_HTML = """
<td><a class="external text"
  data-mw='{"parts":[{"template":{"target":{"wt":"NyseSymbol","href":"./Template:NyseSymbol"},"params":{"1":{"wt":"BNY"}},"i":0}}]}'
  >BNY</a></td>
<td><a href="https://www.nyse.com/quote/XNYS:DVN">DVN</a></td>
<td><a href="https://www.nasdaq.com/market-activity/stocks/aapl">AAPL</a></td>
<td><a href="https://www.nyse.com/quote/XNYS:MMM">MMM</a></td>
"""


def test_extract_wiki_symbols_finds_bny_and_exchange_links():
    raw = _extract_wiki_symbols(_SAMPLE_WIKI_HTML)
    symbols = set(_normalise_symbols(raw))
    assert "BNY" in symbols
    assert "DVN" in symbols
    assert "AAPL" in symbols
    assert "MMM" in symbols
    assert "BK" not in symbols


def test_legacy_td_anchor_still_works():
    html = '<tr><td><a href="/wiki/foo">BK</a></td></tr>'
    symbols = set(_normalise_symbols(_extract_wiki_symbols(html)))
    assert "BK" in symbols