chinese-equity-quant/tests/test_joinquant_plugin.py

"""Tests for the JoinQuant comparison plugin (network-free)."""

from __future__ import annotations

import hashlib
import json
from pathlib import Path

import pandas as pd
import pytest
from click.testing import CliRunner

from cli import cli
from pipeline.common.schema import FILL_COLUMNS, PNL_COLUMNS, POSITION_COLUMNS
from plugins.joinquant.export_targets import export_targets
from plugins.joinquant.ingest import (
    ingest_joinquant_outputs,
    normalize_fills_csv,
)
from plugins.joinquant.reconcile import reconcile_joinquant
from plugins.joinquant.schema import (
    JOINQUANT_FILL_COLUMNS,
    JOINQUANT_PNL_COLUMNS,
    JOINQUANT_POSITION_COLUMNS,
    JOINQUANT_TARGET_COLUMNS,
    RECONCILE_COLUMNS,
)
from plugins.joinquant.smoke import build_fixed_share_positions
from plugins.joinquant.symbols import from_joinquant_symbol, to_joinquant_symbol
from plugins.joinquant.wrapper_strategy import write_wrapper_strategy


def _positions(
    *,
    symbol: str = "sh600000",
    date: str = "2026-07-01",
    shares: int = 1000,
    price: float = 10.0,
    portfolio_name: str = "run1",
) -> pd.DataFrame:
    target_value = float(shares * price)
    weight = target_value / 1_000_000.0
    return pd.DataFrame([{
        "symbol_id": symbol,
        "date": pd.Timestamp(date),
        "portfolio_name": portfolio_name,
        "target_weight": weight,
        "target_value": target_value,
        "target_shares": float(shares) + 0.25,
        "position_shares": shares,
        "position_value": target_value,
        "price": price,
    }], columns=POSITION_COLUMNS)


def _our_fills(
    *,
    symbol: str = "sh600000",
    date: str = "2026-07-01",
    shares: int = 1000,
    price: float = 10.0,
    cost: float = 5.0,
    portfolio_name: str = "run1",
) -> pd.DataFrame:
    fills = pd.DataFrame([{
        "symbol_id": symbol,
        "date": pd.Timestamp(date),
        "portfolio_name": portfolio_name,
        "prev_shares": 0,
        "target_shares": shares,
        "traded_shares": shares,
        "realized_shares": shares,
        "blocked": 0,
        "trade_cost": cost,
        "trade_price": price,
    }])
    return fills


def _our_pnl(
    *,
    date: str = "2026-07-01",
    pnl: float = 100.0,
    cost: float = 5.0,
    portfolio_name: str = "run1",
) -> pd.DataFrame:
    return pd.DataFrame([{
        "date": pd.Timestamp(date),
        "portfolio_name": portfolio_name,
        "gross_exposure": 10_000.0,
        "net_exposure": 10_000.0,
        "pnl": pnl,
        "cost": cost,
        "turnover": 1.0,
        "n_positions": 1,
    }], columns=PNL_COLUMNS)


def _jq_fills(
    *,
    symbol: str = "sh600000",
    date: str = "2026-07-01",
    shares: int = 1000,
    price: float = 10.0,
    cost: float = 5.0,
    portfolio_name: str = "run1",
    raw_status: str = "filled",
) -> pd.DataFrame:
    return pd.DataFrame([{
        "date": date,
        "portfolio_name": portfolio_name,
        "symbol_id": symbol,
        "jq_symbol": to_joinquant_symbol(symbol),
        "order_id": "ord-1",
        "side": "buy" if shares >= 0 else "sell",
        "requested_shares": shares,
        "filled_shares": shares,
        "fill_price": price,
        "trade_value": abs(shares * price),
        "trade_cost": cost,
        "blocked": 0,
        "raw_status": raw_status,
    }], columns=JOINQUANT_FILL_COLUMNS)


def _jq_positions(
    *,
    symbol: str = "sh600000",
    date: str = "2026-07-01",
    shares: int = 1000,
    price: float = 10.0,
    portfolio_name: str = "run1",
) -> pd.DataFrame:
    return pd.DataFrame([{
        "date": date,
        "portfolio_name": portfolio_name,
        "symbol_id": symbol,
        "jq_symbol": to_joinquant_symbol(symbol),
        "position_shares": shares,
        "position_value": shares * price,
        "cash": 990_000.0,
        "total_value": 1_000_000.0,
    }], columns=JOINQUANT_POSITION_COLUMNS)


def _jq_pnl(
    *,
    date: str = "2026-07-01",
    pnl: float = 100.0,
    cost: float = 5.0,
    portfolio_name: str = "run1",
) -> pd.DataFrame:
    return pd.DataFrame([{
        "date": date,
        "portfolio_name": portfolio_name,
        "gross_exposure": 10_000.0,
        "net_exposure": 10_000.0,
        "cash": 990_000.0,
        "total_value": 1_000_000.0,
        "pnl": pnl,
        "cost": cost,
        "turnover": 1.0,
    }], columns=JOINQUANT_PNL_COLUMNS)


def _write_parquets(tmp_path: Path, frames: dict[str, pd.DataFrame]) -> dict[str, Path]:
    paths = {}
    for name, frame in frames.items():
        path = tmp_path / f"{name}.pq"
        frame.to_parquet(path, index=False)
        paths[name] = path
    return paths


def _export_targets_for(tmp_path: Path, positions: pd.DataFrame) -> tuple[Path, Path]:
    positions_path = tmp_path / "positions.pq"
    positions.to_parquet(positions_path, index=False)
    targets_root = tmp_path / "targets"
    export_targets(
        positions_path,
        portfolio_name="run1",
        out_dir=targets_root,
        mode="target_shares",
    )
    return positions_path, targets_root / "run1"


@pytest.mark.parametrize(
    ("internal", "joinquant"),
    [
        ("sh600000", "600000.XSHG"),
        ("sh688001", "688001.XSHG"),
        ("sz000001", "000001.XSHE"),
        ("sz001001", "001001.XSHE"),
        ("sz002594", "002594.XSHE"),
        ("sz300001", "300001.XSHE"),
    ],
)
def test_symbol_mapping_both_directions(internal, joinquant):
    assert to_joinquant_symbol(internal) == joinquant
    assert from_joinquant_symbol(joinquant) == internal


@pytest.mark.parametrize("bad", ["600000", "bj830000", "sh000001", "sz600000", "abc"])
def test_symbol_mapping_rejects_invalid_symbols(bad):
    with pytest.raises(ValueError):
        to_joinquant_symbol(bad)


@pytest.mark.parametrize("bad", ["600000", "600000.XSHE", "000001.XSHG", "abc.XSHG"])
def test_reverse_symbol_mapping_rejects_invalid_symbols(bad):
    with pytest.raises(ValueError):
        from_joinquant_symbol(bad)


def test_export_targets_schema_snapshot_hash_and_no_overwrite(tmp_path):
    positions_path = tmp_path / "positions.pq"
    _positions().to_parquet(positions_path, index=False)

    snapshots = export_targets(
        positions_path,
        portfolio_name="run1",
        out_dir=tmp_path / "targets",
        mode="target_shares",
    )

    csv_path = tmp_path / "targets" / "run1" / "20260701.csv"
    parquet_path = tmp_path / "targets" / "run1" / "20260701.parquet"
    snapshot_path = tmp_path / "snapshots" / "run1" / "20260701.json"
    assert csv_path.exists()
    assert parquet_path.exists()
    assert snapshot_path.exists()

    target = pd.read_csv(csv_path)
    assert list(target.columns) == JOINQUANT_TARGET_COLUMNS
    assert int(target.loc[0, "target_shares"]) == 1000
    assert float(target.loc[0, "target_value"]) == 10_000.0
    assert target.loc[0, "export_mode"] == "target_shares"

    snapshot = json.loads(snapshot_path.read_text())
    actual_hash = hashlib.sha256(csv_path.read_bytes()).hexdigest()
    assert snapshots[0]["file_sha256"] == actual_hash
    assert snapshot["file_sha256"] == actual_hash
    assert snapshot["n_symbols"] == 1

    with pytest.raises(FileExistsError):
        export_targets(
            positions_path,
            portfolio_name="run1",
            out_dir=tmp_path / "targets",
            mode="target_shares",
        )


def test_export_targets_target_value_mode_from_position_columns(tmp_path):
    positions_path = tmp_path / "positions.pq"
    _positions(shares=250, price=20.0).to_parquet(positions_path, index=False)

    export_targets(
        positions_path,
        portfolio_name="run1",
        out_dir=tmp_path / "targets_value",
        mode="target_value",
    )

    target = pd.read_parquet(tmp_path / "targets_value" / "run1" / "20260701.parquet")
    assert list(target.columns) == JOINQUANT_TARGET_COLUMNS
    assert target.loc[0, "export_mode"] == "target_value"
    assert target.loc[0, "target_value"] == 5_000.0
    assert target.loc[0, "target_shares"] == 250


def test_export_targets_can_shift_to_next_execution_session(tmp_path):
    positions_path = tmp_path / "positions.pq"
    _positions(date="2024-01-09").to_parquet(positions_path, index=False)
    calendar_path = tmp_path / "daily.pq"
    pd.DataFrame({
        "date": pd.to_datetime(["2024-01-09", "2024-01-10", "2024-01-11"]),
        "symbol_id": ["sh600000", "sh600000", "sh600000"],
    }).to_parquet(calendar_path, index=False)

    snapshots = export_targets(
        positions_path,
        portfolio_name="run1",
        out_dir=tmp_path / "targets_shifted",
        mode="target_shares",
        start_date="2024-01-10",
        end_date="2024-01-10",
        execution_calendar_path=calendar_path,
    )

    assert len(snapshots) == 1
    assert snapshots[0]["date"] == "2024-01-10"
    assert (tmp_path / "targets_shifted" / "run1" / "20240110.csv").exists()
    target = pd.read_csv(tmp_path / "targets_shifted" / "run1" / "20240110.csv")
    assert target.loc[0, "date"] == "2024-01-10"


def test_ingest_permissive_csv_column_mapping_and_output_schemas(tmp_path):
    fills_csv = tmp_path / "jq_fills.csv"
    positions_csv = tmp_path / "jq_positions.csv"
    pnl_csv = tmp_path / "jq_pnl.csv"
    pd.DataFrame([{
        "Trade Date": "2026-07-01 09:31:00",
        "Security": "600000.XSHG",
        "Direction": "buy",
        "Order Amount": 1000,
        "Filled Amount": 1000,
        "Price": 10.0,
        "Status": "filled",
    }]).to_csv(fills_csv, index=False)
    pd.DataFrame([{
        "Date": "2026-07-01",
        "Security": "600000.XSHG",
        "Shares": 1000,
        "Market Value": 10_000.0,
        "Cash": 990_000.0,
        "Portfolio Value": 1_000_000.0,
    }]).to_csv(positions_csv, index=False)
    pd.DataFrame([{
        "Date": "2026-07-01",
        "Portfolio Value": 1_000_000.0,
        "Daily PnL": 100.0,
        "Turnover": 1.0,
    }]).to_csv(pnl_csv, index=False)

    fills = normalize_fills_csv(fills_csv, "run1")
    assert list(fills.columns) == JOINQUANT_FILL_COLUMNS
    assert fills.loc[0, "symbol_id"] == "sh600000"
    assert fills.loc[0, "jq_symbol"] == "600000.XSHG"
    assert fills.loc[0, "trade_cost"] == 0.0
    assert fills.loc[0, "blocked"] == 0

    paths = ingest_joinquant_outputs(
        portfolio_name="run1",
        fills_csv=fills_csv,
        positions_csv=positions_csv,
        pnl_csv=pnl_csv,
        out_dir=tmp_path / "ingested",
    )
    assert list(pd.read_parquet(paths["fills"]).columns) == JOINQUANT_FILL_COLUMNS
    assert list(pd.read_parquet(paths["positions"]).columns) == JOINQUANT_POSITION_COLUMNS
    assert list(pd.read_parquet(paths["pnl"]).columns) == JOINQUANT_PNL_COLUMNS


def _run_reconcile_case(
    tmp_path: Path,
    *,
    positions: pd.DataFrame | None = None,
    our_fills: pd.DataFrame | None = None,
    jq_fills: pd.DataFrame | None = None,
    jq_positions: pd.DataFrame | None = None,
    our_pnl: pd.DataFrame | None = None,
    jq_pnl: pd.DataFrame | None = None,
) -> pd.DataFrame:
    positions = _positions() if positions is None else positions
    _, targets_dir = _export_targets_for(tmp_path, positions)
    paths = _write_parquets(tmp_path, {
        "our_fills": _our_fills() if our_fills is None else our_fills,
        "our_positions": positions,
        "our_pnl": _our_pnl() if our_pnl is None else our_pnl,
        "jq_fills": _jq_fills() if jq_fills is None else jq_fills,
        "jq_positions": _jq_positions() if jq_positions is None else jq_positions,
        "jq_pnl": _jq_pnl() if jq_pnl is None else jq_pnl,
    })
    out_paths = reconcile_joinquant(
        portfolio_name="run1",
        targets_dir=targets_dir,
        our_fills_path=paths["our_fills"],
        our_positions_path=paths["our_positions"],
        our_pnl_path=paths["our_pnl"],
        jq_fills_path=paths["jq_fills"],
        jq_positions_path=paths["jq_positions"],
        jq_pnl_path=paths["jq_pnl"],
        out_dir=tmp_path / "reconcile",
    )
    report = pd.read_parquet(out_paths["daily_reconcile"])
    assert list(report.columns) == RECONCILE_COLUMNS
    assert out_paths["summary_md"].exists()
    assert out_paths["summary_csv"].exists()
    return report


def test_reconcile_exact_match(tmp_path):
    report = _run_reconcile_case(tmp_path)
    assert report.loc[0, "diff_reason"] == "MATCH"
    assert report.loc[0, "filled_share_diff"] == 0
    assert report.loc[0, "position_share_diff"] == 0


def test_reconcile_price_mismatch(tmp_path):
    report = _run_reconcile_case(tmp_path, jq_fills=_jq_fills(price=10.5))
    assert report.loc[0, "diff_reason"] == "PRICE_MISMATCH"


def test_reconcile_cost_mismatch(tmp_path):
    report = _run_reconcile_case(
        tmp_path,
        jq_fills=_jq_fills(cost=8.0),
        jq_pnl=_jq_pnl(cost=8.0),
    )
    assert report.loc[0, "diff_reason"] == "COST_MODEL"


def test_reconcile_missing_symbol_in_joinquant(tmp_path):
    empty_jq_fills = pd.DataFrame(columns=JOINQUANT_FILL_COLUMNS)
    empty_jq_positions = pd.DataFrame(columns=JOINQUANT_POSITION_COLUMNS)
    report = _run_reconcile_case(
        tmp_path,
        jq_fills=empty_jq_fills,
        jq_positions=empty_jq_positions,
    )
    assert report.loc[0, "diff_reason"] == "MISSING_IN_JOINQUANT"


def test_reconcile_short_target_with_long_only_joinquant_output(tmp_path):
    positions = _positions(shares=-100, price=10.0)
    our_fills = _our_fills(shares=-100, price=10.0)
    jq_fills = _jq_fills(shares=0, price=10.0, cost=0.0, raw_status="short clipped")
    jq_positions = _jq_positions(shares=0, price=10.0)

    report = _run_reconcile_case(
        tmp_path,
        positions=positions,
        our_fills=our_fills,
        jq_fills=jq_fills,
        jq_positions=jq_positions,
    )
    assert report.loc[0, "diff_reason"] == "SHORT_NOT_SUPPORTED"


def test_joinquant_cli_smoke_export_ingest_reconcile_and_wrapper(tmp_path):
    runner = CliRunner()
    positions_path = tmp_path / "positions.pq"
    _positions().to_parquet(positions_path, index=False)

    result = runner.invoke(cli, [
        "joinquant", "export-targets",
        "--positions-path", str(positions_path),
        "--portfolio-name", "run1",
        "--mode", "target_shares",
        "--out-dir", str(tmp_path / "targets"),
    ])
    assert result.exit_code == 0, result.output
    assert "Exported JoinQuant targets" in result.output

    fills_csv = tmp_path / "jq_fills.csv"
    positions_csv = tmp_path / "jq_positions.csv"
    pnl_csv = tmp_path / "jq_pnl.csv"
    _jq_fills().to_csv(fills_csv, index=False)
    _jq_positions().to_csv(positions_csv, index=False)
    _jq_pnl().to_csv(pnl_csv, index=False)

    result = runner.invoke(cli, [
        "joinquant", "ingest",
        "--portfolio-name", "run1",
        "--fills-csv", str(fills_csv),
        "--positions-csv", str(positions_csv),
        "--pnl-csv", str(pnl_csv),
        "--out-dir", str(tmp_path / "ingested"),
    ])
    assert result.exit_code == 0, result.output
    assert "Saved JoinQuant fills" in result.output

    paths = _write_parquets(tmp_path, {
        "our_fills": _our_fills(),
        "our_pnl": _our_pnl(),
    })
    result = runner.invoke(cli, [
        "joinquant", "reconcile",
        "--portfolio-name", "run1",
        "--targets-dir", str(tmp_path / "targets" / "run1"),
        "--our-fills-path", str(paths["our_fills"]),
        "--our-positions-path", str(positions_path),
        "--our-pnl-path", str(paths["our_pnl"]),
        "--jq-fills-path", str(tmp_path / "ingested" / "run1" / "fills.pq"),
        "--jq-positions-path", str(tmp_path / "ingested" / "run1" / "positions.pq"),
        "--jq-pnl-path", str(tmp_path / "ingested" / "run1" / "pnl.pq"),
        "--out-dir", str(tmp_path / "reconcile"),
    ])
    assert result.exit_code == 0, result.output
    assert "Saved reconciliation parquet" in result.output

    wrapper_path = tmp_path / "wrapper_strategy_run1.py"
    result = runner.invoke(cli, [
        "joinquant", "write-wrapper",
        "--portfolio-name", "run1",
        "--mode", "target_shares",
        "--out-path", str(wrapper_path),
    ])
    assert result.exit_code == 0, result.output
    assert "Saved JoinQuant wrapper strategy" in result.output
    text = wrapper_path.read_text()
    assert 'PORTFOLIO_NAME = "run1"' in text
    assert 'TARGET_MODE = "target_shares"' in text
    assert "ALLOW_SHORT = False" in text


def test_wrapper_strategy_generation_smoke(tmp_path):
    path = write_wrapper_strategy(
        portfolio_name="run2",
        mode="target_value",
        out_path=tmp_path / "wrapper.py",
    )
    text = path.read_text()
    assert 'PORTFOLIO_NAME = "run2"' in text
    assert 'TARGET_MODE = "target_value"' in text
    assert "order_target_value" in text


def test_build_fixed_share_positions_excludes_final_executionless_date():
    data = pd.DataFrame({
        "symbol_id": ["sh600000", "sh600000", "sh600000"],
        "date": pd.to_datetime(["2024-01-09", "2024-01-10", "2024-01-11"]),
        "close": [10.0, 10.5, 11.0],
    })

    positions = build_fixed_share_positions(
        data,
        trade_symbol="sh600000",
        portfolio_name="run1",
        shares=1000,
        booksize=1_000_000.0,
    )

    assert list(positions.columns) == POSITION_COLUMNS
    assert positions["date"].dt.strftime("%Y-%m-%d").tolist() == [
        "2024-01-09",
        "2024-01-10",
    ]
    assert positions["position_shares"].tolist() == [1000, 1000]
    assert positions["target_value"].tolist() == [10_000.0, 10_500.0]