zhixing-system/zhixing-server/tests/unit/selection/test_chart.py

"""Application tests for bounded qfq chart series."""

from datetime import date, timedelta

import pandas as pd
import pytest

from zhixing_server.modules.selection.application.chart import (
    GetSelectionChart,
    SelectionChartNotFound,
)
from zhixing_server.modules.selection.domain.gold_brick import (
    prepare_gold_brick_indicators,
)
from zhixing_server.modules.selection.domain.indicators import (
    compute_kdj,
    compute_zhixing_lines,
)
from zhixing_server.modules.selection.domain.models import SelectionBar, StockHistory


class FakeReader:
    """Return one deterministic history through the requested date."""

    def __init__(self, history: StockHistory) -> None:
        self.history = history
        self.request: tuple[str, date] | None = None

    def load_history(self, ts_code: str, target_trade_date: date) -> StockHistory:
        self.request = (ts_code, target_trade_date)
        return self.history


def _history(count: int) -> StockHistory:
    start = date(2025, 1, 1)
    bars = tuple(
        SelectionBar(
            trade_date=start + timedelta(days=index),
            open=10.0 + index / 100,
            high=11.0 + index / 100,
            low=9.0 + index / 100,
            close=10.5 + index / 100,
            volume=1000.0 + index,
        )
        for index in range(count)
    )
    return StockHistory(ts_code="000001.SZ", name="平安银行", bars=bars)


def test_chart_computes_on_full_history_before_returning_last_250_points() -> None:
    history = _history(260)
    reader = FakeReader(history)
    target = history.bars[-1].trade_date

    chart = GetSelectionChart(reader).execute(history.ts_code, target)

    assert reader.request == ("000001.SZ", target)
    assert chart.source_adj == "qfq"
    assert len(chart.points) == 250
    assert chart.points[0].trade_date == history.bars[10].trade_date
    assert chart.points[-1].trade_date == target
    frame = pd.DataFrame(
        {
            "low": [bar.low for bar in history.bars],
            "high": [bar.high for bar in history.bars],
            "close": [bar.close for bar in history.bars],
        }
    )
    expected = compute_kdj(frame).iloc[-1]
    expected_white, expected_yellow = compute_zhixing_lines(frame["close"])
    assert chart.points[-1].k == pytest.approx(float(expected["K"]))
    assert chart.points[-1].d == pytest.approx(float(expected["D"]))
    assert chart.points[-1].j == pytest.approx(float(expected["J"]))
    assert chart.points[-1].trend_white == pytest.approx(float(expected_white.iloc[-1]))
    assert chart.points[-1].trend_yellow == pytest.approx(float(expected_yellow.iloc[-1]))


def test_chart_filters_future_rows_and_preserves_nullable_points() -> None:
    history = _history(12)
    target = history.bars[-2].trade_date
    incomplete = SelectionBar(
        trade_date=history.bars[3].trade_date,
        open=history.bars[3].open,
        high=None,
        low=history.bars[3].low,
        close=None,
        volume=None,
    )
    history = StockHistory(
        ts_code=history.ts_code,
        name=history.name,
        bars=history.bars[:3] + (incomplete,) + history.bars[4:],
    )

    chart = GetSelectionChart(FakeReader(history)).execute(history.ts_code, target)

    assert chart.points[-1].trade_date == target
    assert all(point.trade_date <= target for point in chart.points)
    assert chart.points[3].high is None
    assert chart.points[3].volume is None
    assert chart.points[3].k is None
    assert chart.points[3].d is None
    assert chart.points[3].j is None


def test_chart_skips_brick_series_by_default_and_computes_on_request() -> None:
    history = _history(260)
    reader = FakeReader(history)
    target = history.bars[-1].trade_date

    default_chart = GetSelectionChart(reader).execute(history.ts_code, target)
    assert all(point.brick_chart is None for point in default_chart.points)

    brick_chart = GetSelectionChart(reader).execute(
        history.ts_code,
        target,
        include_brick_chart=True,
    )

    frame = pd.DataFrame(
        {
            "trade_date": [bar.trade_date for bar in history.bars],
            "open": [bar.open for bar in history.bars],
            "high": [bar.high for bar in history.bars],
            "low": [bar.low for bar in history.bars],
            "close": [bar.close for bar in history.bars],
            "volume": [bar.volume for bar in history.bars],
        }
    )
    prepared = prepare_gold_brick_indicators(frame, history.ts_code)
    assert brick_chart.points[-1].brick_chart == pytest.approx(
        float(prepared["brick_chart"].iloc[-1])
    )
    assert brick_chart.points[-1].brick_strong_red is bool(prepared["strong_red"].iloc[-1])
    assert brick_chart.points[0].trade_date == history.bars[10].trade_date
    assert brick_chart.points[0].brick_chart == pytest.approx(
        float(prepared["brick_chart"].iloc[10])
    )


def test_chart_rejects_empty_history() -> None:
    reader = FakeReader(StockHistory(ts_code="000001.SZ", name="平安银行"))

    with pytest.raises(SelectionChartNotFound, match="chart data not found"):
        GetSelectionChart(reader).execute("000001.SZ", date(2026, 8, 8))