crypto_monitor/tests/test_hub_kline_store.py

"""中控 K 线库：分周期保留、聚合与分页读取。"""
from __future__ import annotations

import tempfile
import time
import unittest
from pathlib import Path

from hub_kline_store import (
    init_db,
    load_bars_before,
    load_bars_latest,
    purge_retention,
    purge_timeframe_by_days,
    resolve_chart_bars,
    retention_days,
    upsert_bars,
)
from hub_ohlcv_lib import (
    TIMEFRAME_MS,
    bar_limit_for_timeframe,
    chart_fetch_start_ms,
    chart_initial_limit,
    last_closed_bar_open_ms,
    window_start_ms,
)


class TestHubKlineStore(unittest.TestCase):
    def setUp(self):
        self.tmp = tempfile.TemporaryDirectory()
        self.db = Path(self.tmp.name) / "test_hub_kline.db"

    def tearDown(self):
        self.tmp.cleanup()

    def test_bar_limits(self):
        self.assertEqual(bar_limit_for_timeframe("5m"), 5000)
        self.assertEqual(bar_limit_for_timeframe("1h"), 1000)
        self.assertEqual(bar_limit_for_timeframe("1d"), 1000)
        self.assertEqual(bar_limit_for_timeframe("1w"), 500)
        self.assertEqual(chart_initial_limit("5m"), 300)

    def test_chart_fetch_window_exceeds_retention(self):
        now = int(time.time() * 1000)
        need = bar_limit_for_timeframe("1d")
        fetch_start = chart_fetch_start_ms("1d", need, now)
        db_start = window_start_ms("1d", need, retention_days(), now)
        self.assertLess(fetch_start, db_start)

    def test_purge_retention_5m_one_year(self):
        init_db(self.db)
        old_ms = int(time.time() * 1000) - 400 * 86400000
        upsert_bars(
            "okx",
            "BTC/USDT",
            "5m",
            [
                {
                    "open_time_ms": old_ms,
                    "open": 1,
                    "high": 2,
                    "low": 0.5,
                    "close": 1.5,
                    "volume": 10,
                }
            ],
            self.db,
        )
        n = purge_timeframe_by_days("5m", 365, self.db)
        self.assertGreaterEqual(n, 1)
        rows = load_bars_latest("okx", "BTC/USDT", "5m", 10, self.db)
        self.assertEqual(len(rows), 0)

    def test_purge_retention_keeps_1d(self):
        init_db(self.db)
        old_ms = int(time.time() * 1000) - 400 * 86400000
        upsert_bars(
            "okx",
            "BTC/USDT",
            "1d",
            [
                {
                    "open_time_ms": old_ms,
                    "open": 1,
                    "high": 2,
                    "low": 0.5,
                    "close": 1.5,
                    "volume": 10,
                }
            ],
            self.db,
        )
        purge_retention(self.db)
        rows = load_bars_latest("okx", "BTC/USDT", "1d", 10, self.db)
        self.assertEqual(len(rows), 1)

    def test_resolve_uses_cache_without_remote(self):
        init_db(self.db)
        now = int(time.time() * 1000)
        tf = "5m"
        period = TIMEFRAME_MS[tf]
        last_closed = last_closed_bar_open_ms(tf, now)
        bars = []
        for i in range(400):
            oms = last_closed - (399 - i) * period
            bars.append(
                {
                    "open_time_ms": oms,
                    "open": 100 + i,
                    "high": 101 + i,
                    "low": 99 + i,
                    "close": 100.5 + i,
                    "volume": 1000 + i,
                }
            )
        upsert_bars("okx", "ETH/USDT", tf, bars, self.db)

        def remote_fetch(**kwargs):
            self.fail("不应请求交易所")

        out = resolve_chart_bars(
            "okx",
            "ETH/USDT",
            tf,
            remote_fetch,
            db_path=self.db,
            limit=300,
        )
        self.assertTrue(out.get("ok"))
        self.assertEqual(len(out.get("candles") or []), 300)

    def test_resolve_15m_reads_native_bars(self):
        init_db(self.db)
        now = int(time.time() * 1000)
        period = TIMEFRAME_MS["15m"]
        last_closed = last_closed_bar_open_ms("15m", now)
        bars = []
        for i in range(12):
            oms = last_closed - (11 - i) * period
            bars.append(
                {
                    "open_time_ms": oms,
                    "open": 1.0 + i,
                    "high": 2.0 + i,
                    "low": 0.5 + i,
                    "close": 1.5 + i,
                    "volume": 10.0,
                }
            )
        upsert_bars("okx", "ETH/USDT", "15m", bars, self.db)

        def remote_fetch(**kwargs):
            self.fail("不应请求交易所")

        out = resolve_chart_bars(
            "okx",
            "ETH/USDT",
            "15m",
            remote_fetch,
            db_path=self.db,
            limit=10,
        )
        self.assertTrue(out.get("ok"))
        self.assertEqual(out.get("source"), "db")
        self.assertEqual(out.get("storage_timeframe"), "15m")
        self.assertGreaterEqual(len(out.get("candles") or []), 10)

    def test_load_bars_before(self):
        init_db(self.db)
        period = TIMEFRAME_MS["1h"]
        base = 1_700_000_000_000
        bars = []
        for i in range(5):
            bars.append(
                {
                    "open_time_ms": base + i * period,
                    "open": 1,
                    "high": 2,
                    "low": 0.5,
                    "close": 1.5,
                    "volume": 1,
                }
            )
        upsert_bars("okx", "BTC/USDT", "1h", bars, self.db)
        before = base + 3 * period
        got = load_bars_before("okx", "BTC/USDT", "1h", before, 2, self.db)
        self.assertEqual(len(got), 2)
        self.assertEqual(got[-1]["open_time_ms"], base + 2 * period)

    def test_resolve_before_ms_exhausted(self):
        init_db(self.db)

        def remote_fetch(**kwargs):
            return {"ok": False, "msg": "no remote"}

        out = resolve_chart_bars(
            "okx",
            "BTC/USDT",
            "5m",
            remote_fetch,
            db_path=self.db,
            limit=100,
            before_ms=int(time.time() * 1000),
        )
        self.assertTrue(out.get("ok"))
        self.assertEqual(out.get("candles"), [])
        self.assertTrue(out.get("exhausted"))


if __name__ == "__main__":
    unittest.main()