90 lines
3.5 KiB
Python
90 lines
3.5 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()
|
|
|
|
def test_realtime_warmup_runs_once_per_interval_during_market(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,
|
|
quality={"realtime_warmup_interval_seconds": 120},
|
|
)
|
|
pipe = Pipeline(db, adapter, settings)
|
|
pipe.ingest_reference("20240902")
|
|
calls: list[str] = []
|
|
sched = Scheduler(db, pipe, jobs={"realtime_warmup": lambda day: calls.append(day) or {}})
|
|
|
|
first = sched.tick(datetime(2024, 9, 2, 10, 0, 0, tzinfo=SHANGHAI))
|
|
duplicate = sched.tick(datetime(2024, 9, 2, 10, 1, 0, tzinfo=SHANGHAI))
|
|
second = sched.tick(datetime(2024, 9, 2, 10, 2, 0, tzinfo=SHANGHAI))
|
|
lunch = sched.tick(datetime(2024, 9, 2, 12, 0, 0, tzinfo=SHANGHAI))
|
|
|
|
self.assertIn("realtime_warmup", first)
|
|
self.assertNotIn("realtime_warmup", duplicate)
|
|
self.assertIn("realtime_warmup", second)
|
|
self.assertNotIn("realtime_warmup", lunch)
|
|
self.assertEqual(calls, ["20240902", "20240902"])
|
|
tmp.cleanup()
|
|
|
|
|
|
if __name__ == "__main__":
|
|
unittest.main()
|