tests/test_external_factor_weight_grid.py

import pandas as pd
import pytest

from scripts.search_external_factor_weight_grid import (
    load_external_factor_scores,
    merge_external_scores,
)


def test_load_external_factor_scores_normalizes_and_combines_sources(tmp_path):
    chipk = tmp_path / "chipk.csv"
    crowd = tmp_path / "crowd.csv"
    chipk.write_text(
        "stock,date,bull_score,bear_score\n"
        "2330,2026/05/01,2,0\n"
        "50,2026-05-01,1,0\n",
        encoding="utf-8",
    )
    crowd.write_text(
        "stock,date,bull_score,bear_score\n"
        "2330.TW,2026-05-01,0.5,1\n",
        encoding="utf-8",
    )

    scores = load_external_factor_scores([chipk, crowd])

    row_2330 = scores[(scores["stock"] == "2330") & (scores["date"] == "2026-05-01")].iloc[0]
    row_0050 = scores[(scores["stock"] == "0050") & (scores["date"] == "2026-05-01")].iloc[0]
    assert row_2330["bull_score"] == 2.5
    assert row_2330["bear_score"] == 1.0
    assert row_2330["source_count"] == 2
    assert row_0050["bull_score"] == 1.0


def test_load_external_factor_scores_requires_configured_columns(tmp_path):
    path = tmp_path / "bad.csv"
    path.write_text("stock,date,bull_score\n2330,2026-05-01,1\n", encoding="utf-8")

    with pytest.raises(ValueError, match="missing required column"):
        load_external_factor_scores([path])


def test_merge_external_scores_replaces_generic_factor_scores():
    events = pd.DataFrame(
        {
            "stock": ["2330", "2330", "0050"],
            "date": ["2026-05-01", "2026-05-02", "2026-05-01"],
            "p_buy": [0.4, 0.4, 0.4],
            "p_hold": [0.5, 0.5, 0.5],
            "p_sell": [0.1, 0.1, 0.1],
            "y_true": [1, 0, 1],
            "base_pred": [0, 0, 0],
            "bull_score": [9.0, 9.0, 9.0],
            "bear_score": [9.0, 9.0, 9.0],
        }
    )
    external = pd.DataFrame(
        {
            "stock": ["2330"],
            "date": ["2026-05-01"],
            "bull_score": [3.0],
            "bear_score": [1.0],
            "source_count": [1],
        }
    )

    merged, summary = merge_external_scores(events, external)

    assert summary["matched_event_count"] == 1
    assert summary["coverage"] == 0.3333
    assert merged.loc[0, "bull_score"] == 3.0
    assert merged.loc[0, "bear_score"] == 1.0
    assert merged.loc[1, "bull_score"] == 0.0
    assert merged.loc[2, "bear_score"] == 0.0