3498dd7a4b
新增独立 xiaobai-datahub 服务（SQLite WAL、Tushare 盘后发布、/v1 契约和管理后台），不改现站页面与数据链路。 Co-authored-by: Cursor <cursoragent@cursor.com> Co-authored-by: multica-agent <github@multica.ai>
63 lines
2.3 KiB
Python
63 lines
2.3 KiB
Python
from __future__ import annotations
|
|
|
|
import unittest
|
|
from datetime import datetime
|
|
from pathlib import Path
|
|
import tempfile
|
|
|
|
from datahub.adapters.tushare import TushareAdapter
|
|
from datahub.crypto import SecretVault
|
|
from datahub.db import HubDB
|
|
from datahub.pipeline import Pipeline
|
|
from datahub.scheduler import Scheduler
|
|
from datahub.settings import Settings
|
|
from datahub.timeutil import SHANGHAI
|
|
from tests.fixtures import fake_transport
|
|
|
|
|
|
class SchedulerTests(unittest.TestCase):
|
|
def test_skips_eod_on_closed_day(self) -> None:
|
|
tmp = tempfile.TemporaryDirectory()
|
|
db = HubDB(Path(tmp.name) / "hub.db")
|
|
adapter = TushareAdapter("x", transport=fake_transport)
|
|
settings = Settings(encryption_key=SecretVault.generate_key(), scheduler_enabled=False, db_path=db.path)
|
|
pipe = Pipeline(db, adapter, settings)
|
|
pipe.ingest_reference("20240902")
|
|
# 20240907 is closed in fixture
|
|
ran = {"eod_a": 0}
|
|
|
|
def fake_eod(_date: str):
|
|
ran["eod_a"] += 1
|
|
return {}
|
|
|
|
sched = Scheduler(db, pipe, jobs={"precheck": lambda d: {}, "eod_a": fake_eod, "eod_b": lambda d: {}, "cleanup": lambda d: {}, "backup": lambda d: {}})
|
|
clock = datetime(2024, 9, 7, 16, 0, tzinfo=SHANGHAI)
|
|
fired = sched.tick(clock)
|
|
self.assertNotIn("eod_a", fired)
|
|
self.assertEqual(ran["eod_a"], 0)
|
|
tmp.cleanup()
|
|
|
|
def test_fires_eod_on_open_day(self) -> None:
|
|
tmp = tempfile.TemporaryDirectory()
|
|
db = HubDB(Path(tmp.name) / "hub.db")
|
|
adapter = TushareAdapter("x", transport=fake_transport)
|
|
settings = Settings(encryption_key=SecretVault.generate_key(), scheduler_enabled=False, db_path=db.path)
|
|
pipe = Pipeline(db, adapter, settings)
|
|
pipe.ingest_reference("20240902")
|
|
ran = {"eod_a": 0}
|
|
|
|
def fake_eod(_date: str):
|
|
ran["eod_a"] += 1
|
|
return {"rows": 1}
|
|
|
|
sched = Scheduler(db, pipe, jobs={"precheck": lambda d: {}, "eod_a": fake_eod, "eod_b": lambda d: {}, "cleanup": lambda d: {}, "backup": lambda d: {}})
|
|
clock = datetime(2024, 9, 2, 16, 0, tzinfo=SHANGHAI)
|
|
fired = sched.tick(clock)
|
|
self.assertIn("eod_a", fired)
|
|
self.assertEqual(ran["eod_a"], 1)
|
|
tmp.cleanup()
|
|
|
|
|
|
if __name__ == "__main__":
|
|
unittest.main()
|