momentum-system/tests/unit/research/test_news_enrichment.py

"""Local-LLM (Ollama) news/catalyst enrichment tests.

The LLM call is mocked at ``urllib.request.urlopen`` (no network), matching the
codebase's graceful-degradation contract: any failure -> ``None``, never raises.
The persistence tests run against an in-memory Postgres schema.
"""

import io
import json

import pytest

from research.ingestion import news_enrichment as ne
from storage.db import get_connection


# --------------------------------------------------------------------------
# classify_headline / parsing
# --------------------------------------------------------------------------

class _FakeResp:
    def __init__(self, body: bytes):
        self._body = body

    def read(self):
        return self._body

    def __enter__(self):
        return self

    def __exit__(self, *a):
        return False


def _ollama_reply(payload: dict) -> _FakeResp:
    """Wrap a model JSON object the way Ollama's /api/generate returns it."""
    return _FakeResp(json.dumps({"response": json.dumps(payload)}).encode())


def test_classify_headline_happy_path(monkeypatch):
    monkeypatch.setattr(
        ne.urllib.request, "urlopen",
        lambda *a, **k: _ollama_reply({
            "catalyst_type": "fda_approval", "sentiment": 0.9,
            "conviction": 0.85, "is_dilutive": False, "rationale": "FDA nod",
        }),
    )
    a = ne.classify_headline("BioCo gets FDA approval", tickers="BIO")
    assert a is not None
    assert a.catalyst_type == "fda_approval"
    assert a.sentiment == 0.9 and a.conviction == 0.85
    assert a.is_dilutive is False and a.rationale == "FDA nod"


def test_classify_headline_sends_schema_by_default(monkeypatch):
    captured = {}

    def capture(req, *a, **k):
        captured["body"] = json.loads(req.data.decode())
        return _ollama_reply({
            "catalyst_type": "earnings", "sentiment": 0.5,
            "conviction": 0.6, "is_dilutive": False, "rationale": "beat",
        })

    monkeypatch.setattr(ne.urllib.request, "urlopen", capture)
    ne.classify_headline("QRS beats estimates", tickers="QRS")
    fmt = captured["body"]["format"]
    # schema-constrained decoding: format is the JSON schema, enum-locked
    assert isinstance(fmt, dict)
    assert fmt["properties"]["catalyst_type"]["enum"] == sorted(ne.CATALYST_TYPES)
    assert set(fmt["required"]) == {
        "catalyst_type", "sentiment", "conviction", "is_dilutive", "rationale",
    }


def test_classify_headline_use_schema_false_sends_plain_json(monkeypatch):
    captured = {}

    def capture(req, *a, **k):
        captured["body"] = json.loads(req.data.decode())
        return _ollama_reply({
            "catalyst_type": "other", "sentiment": 0.0,
            "conviction": 0.0, "is_dilutive": False, "rationale": "",
        })

    monkeypatch.setattr(ne.urllib.request, "urlopen", capture)
    ne.classify_headline("anything", use_schema=False)
    assert captured["body"]["format"] == "json"


def test_classify_headline_network_down_returns_none(monkeypatch):
    def boom(*a, **k):
        raise OSError("connection refused")
    monkeypatch.setattr(ne.urllib.request, "urlopen", boom)
    assert ne.classify_headline("anything") is None


def test_classify_headline_malformed_json_returns_none(monkeypatch):
    monkeypatch.setattr(
        ne.urllib.request, "urlopen",
        lambda *a, **k: _FakeResp(json.dumps({"response": "not json at all"}).encode()),
    )
    assert ne.classify_headline("headline") is None


def test_classify_headline_empty_input_returns_none():
    assert ne.classify_headline("", "") is None


def test_parse_clamps_and_validates():
    a = ne._parse('{"catalyst_type":"bogus","sentiment":5,"conviction":-1}')
    assert a.catalyst_type == "other"      # unknown enum -> other
    assert a.sentiment == 1.0              # clamped to [-1, 1]
    assert a.conviction == 0.0             # clamped to [0, 1]


def test_parse_offering_type_implies_dilutive():
    a = ne._parse('{"catalyst_type":"offering_dilution","sentiment":-0.5,"conviction":0.7,"is_dilutive":false}')
    assert a.is_dilutive is True           # type forces the flag


def test_parse_keyword_floor_catches_missed_dilution():
    # model said not dilutive, but the headline text obviously is
    a = ne._parse(
        '{"catalyst_type":"other","sentiment":0.0,"conviction":0.3,"is_dilutive":false}',
        fallback_text="Company announces $40M registered direct offering",
    )
    assert a.is_dilutive is True


@pytest.mark.parametrize("text,expected", [
    ("Priced public offering of 5M shares", True),
    ("Announces at-the-market facility", True),
    ("Files Form S-3 shelf registration", True),
    ("Reports record Q3 earnings beat", False),
    ("Signs partnership with Pfizer", False),
])
def test_looks_dilutive(text, expected):
    assert ne.looks_dilutive(text) is expected


def test_catalyst_score():
    assert ne.catalyst_score(None) is None
    assert ne.catalyst_score({"conviction": 0.0, "sentiment": 1.0}) == 0.0
    # bearish catalyst -> low score (sentiment floored at 0 in the multiplier)
    assert ne.catalyst_score({"conviction": 0.9, "sentiment": -0.8}) == pytest.approx(0.45)
    assert ne.catalyst_score({"conviction": 0.9, "sentiment": 0.8}) == pytest.approx(0.81)


# --------------------------------------------------------------------------
# Persistence: enrich_recent_news + catalyst_map (in-memory Postgres)
# --------------------------------------------------------------------------

@pytest.fixture
def con():
    c = get_connection(":memory:")
    yield c
    c.close()


class _Cfg:
    """Minimal stand-in for OllamaConfig (avoids pydantic in the persist tests)."""
    host = "http://localhost:11434"
    model = "test-model"
    timeout_seconds = 5
    temperature = 0.3
    max_tokens = 128
    enrichment_lookback_hours = 12
    enrichment_batch_limit = 50


def _seed_news(con, payload_hash, title, tickers, source="prnewswire"):
    con.execute(
        "INSERT INTO raw_news_items (id, fetch_attempt_id, source, raw_url, raw_title, "
        "raw_published_at, raw_body_snippet, raw_tickers, payload_hash, parser_version, "
        "fetched_at, ingest_run_id) VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, "
        "current_timestamp, ?)",
        [payload_hash, "fa1", source, "http://x/1", title, "", "snippet",
         tickers, payload_hash, "rss-v2", "run1"],
    )


def test_enrich_recent_news_persists_and_dedupes(con, monkeypatch):
    _seed_news(con, "hash_abc", "BioCo gets FDA approval", "BIO")
    fixed = ne.CatalystAnalysis(
        catalyst_type="fda_approval", sentiment=0.8, conviction=0.9,
        is_dilutive=False, rationale="approval",
    )
    monkeypatch.setattr(ne, "classify_headline", lambda *a, **k: fixed)

    res = ne.enrich_recent_news(con, _Cfg())
    assert res["enriched"] == 1 and res["errors"] == 0

    # cache row written
    cache = con.execute(
        "SELECT symbol, catalyst_type, sentiment, conviction FROM news_catalyst_cache"
    ).fetchall()
    assert cache == [("BIO", "fda_approval", 0.8, 0.9)]

    # the previously-dead news_events.sentiment is now populated
    ev = con.execute(
        "SELECT symbol, sentiment, category, is_offering FROM news_events"
    ).fetchall()
    assert ev == [("BIO", 0.8, "fda_approval", False)]

    # second pass enriches nothing (deduped by headline hash)
    res2 = ne.enrich_recent_news(con, _Cfg())
    assert res2["enriched"] == 0


def test_enrich_skips_when_llm_unavailable(con, monkeypatch):
    _seed_news(con, "hash_down", "Some headline", "XYZ")
    monkeypatch.setattr(ne, "classify_headline", lambda *a, **k: None)
    res = ne.enrich_recent_news(con, _Cfg())
    assert res["enriched"] == 0 and res["errors"] == 1
    # cache NOT poisoned — headline stays available for a retry next pass
    assert con.execute("SELECT count(*) FROM news_catalyst_cache").fetchone()[0] == 0


def test_enrich_fans_out_multiple_tickers(con, monkeypatch):
    _seed_news(con, "hash_multi", "MergerCo to acquire TargetCo", "ACQ,TGT")
    fixed = ne.CatalystAnalysis(catalyst_type="ma_acquisition", sentiment=0.6,
                                conviction=0.8, is_dilutive=False, rationale="M&A")
    monkeypatch.setattr(ne, "classify_headline", lambda *a, **k: fixed)
    ne.enrich_recent_news(con, _Cfg())
    syms = sorted(r[0] for r in con.execute(
        "SELECT symbol FROM news_catalyst_cache").fetchall())
    assert syms == ["ACQ", "TGT"]


def test_catalyst_map_latest_per_symbol(con, monkeypatch):
    _seed_news(con, "hash_dil", "BioCo announces registered direct offering", "BIO")
    fixed = ne.CatalystAnalysis(catalyst_type="offering_dilution", sentiment=-0.7,
                                conviction=0.85, is_dilutive=True, rationale="dilution")
    monkeypatch.setattr(ne, "classify_headline", lambda *a, **k: fixed)
    ne.enrich_recent_news(con, _Cfg())

    m = ne.catalyst_map(con)
    assert "BIO" in m
    assert m["BIO"]["is_dilutive"] is True
    assert m["BIO"]["catalyst_type"] == "offering_dilution"