dekun
440d1ecbc9
Keep only the latest contiguous K-line segment, purge isolated stale rows, and backfill when the tail is still shorter than the initial limit. Co-authored-by: Cursor <cursoragent@cursor.com>
302 lines
8.9 KiB
Python
302 lines
8.9 KiB
Python
"""中控 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,
|
|
trim_contiguous_tail,
|
|
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"), 2000)
|
|
self.assertEqual(chart_initial_limit("1h"), 1000)
|
|
self.assertEqual(chart_initial_limit("1d"), 500)
|
|
|
|
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_trim_contiguous_tail_drops_orphan_prefix(self):
|
|
period = TIMEFRAME_MS["15m"]
|
|
base_old = 1_700_000_000_000
|
|
base_new = base_old + period * 500
|
|
bars = []
|
|
for i in range(3):
|
|
bars.append(
|
|
{
|
|
"open_time_ms": base_old + i * period,
|
|
"open": 1,
|
|
"high": 2,
|
|
"low": 0.5,
|
|
"close": 1.5,
|
|
"volume": 1,
|
|
}
|
|
)
|
|
for i in range(5):
|
|
bars.append(
|
|
{
|
|
"open_time_ms": base_new + i * period,
|
|
"open": 2,
|
|
"high": 3,
|
|
"low": 1.5,
|
|
"close": 2.5,
|
|
"volume": 2,
|
|
}
|
|
)
|
|
trimmed, split = trim_contiguous_tail(bars, period)
|
|
self.assertEqual(split, 3)
|
|
self.assertEqual(len(trimmed), 5)
|
|
self.assertEqual(trimmed[0]["open_time_ms"], base_new)
|
|
|
|
def test_resolve_drops_discontinuous_orphans(self):
|
|
init_db(self.db)
|
|
period = TIMEFRAME_MS["15m"]
|
|
now = int(time.time() * 1000)
|
|
old_ms = now - period * 800
|
|
upsert_bars(
|
|
"okx",
|
|
"ONDO/USDT",
|
|
"15m",
|
|
[
|
|
{
|
|
"open_time_ms": old_ms,
|
|
"open": 0.33,
|
|
"high": 0.34,
|
|
"low": 0.32,
|
|
"close": 0.335,
|
|
"volume": 100,
|
|
}
|
|
],
|
|
self.db,
|
|
)
|
|
recent = []
|
|
start = now - period * 20
|
|
for i in range(20):
|
|
recent.append(
|
|
{
|
|
"open_time_ms": start + i * period,
|
|
"open": 0.35,
|
|
"high": 0.36,
|
|
"low": 0.34,
|
|
"close": 0.355,
|
|
"volume": 50,
|
|
}
|
|
)
|
|
|
|
def remote_fetch(**kwargs):
|
|
return {"ok": True, "bars": recent, "price_tick": 0.0001}
|
|
|
|
out = resolve_chart_bars(
|
|
"okx",
|
|
"ONDO/USDT",
|
|
"15m",
|
|
remote_fetch,
|
|
db_path=self.db,
|
|
limit=50,
|
|
)
|
|
self.assertTrue(out.get("ok"))
|
|
candles = out.get("candles") or []
|
|
self.assertGreaterEqual(len(candles), 19)
|
|
if len(candles) >= 2:
|
|
for i in range(1, len(candles)):
|
|
gap = candles[i]["time"] - candles[i - 1]["time"]
|
|
self.assertLessEqual(gap, int(period / 1000 * 1.5))
|
|
|
|
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()
|