419 lines
14 KiB
Python
419 lines
14 KiB
Python
import json
|
|
import tempfile
|
|
import unittest
|
|
from datetime import date, datetime, timezone
|
|
from pathlib import Path
|
|
|
|
from adapters.jqdata_local import (
|
|
JQDataUnavailableError,
|
|
authenticate,
|
|
fetch_index_daily_snapshot,
|
|
)
|
|
from quant60.data_snapshot import (
|
|
CanonicalBarRecord,
|
|
DataQualityError,
|
|
verify_data_snapshot,
|
|
)
|
|
|
|
|
|
class FakeFrame:
|
|
def __init__(self, rows):
|
|
self.rows = rows
|
|
|
|
def iterrows(self):
|
|
return iter(self.rows)
|
|
|
|
|
|
class FakeJQData:
|
|
__version__ = "fake-1"
|
|
|
|
def __init__(self):
|
|
self.authenticated = False
|
|
self.last_get_bars_kwargs = None
|
|
self.last_get_extras_kwargs = None
|
|
self.trade_day_queries = []
|
|
self.full_calendar_queries = 0
|
|
|
|
def auth(self, username, password):
|
|
self.authenticated = username == "user" and password == "password"
|
|
|
|
def is_auth(self):
|
|
return self.authenticated
|
|
|
|
def get_trade_days(self, *, start_date, end_date):
|
|
self.trade_day_queries.append(
|
|
{
|
|
"start_date": start_date,
|
|
"end_date": end_date,
|
|
}
|
|
)
|
|
return [start_date, end_date]
|
|
|
|
def get_all_trade_days(self):
|
|
self.full_calendar_queries += 1
|
|
return [
|
|
date(2024, 1, 2),
|
|
date(2024, 1, 3),
|
|
date(2024, 1, 4),
|
|
]
|
|
|
|
def get_index_stocks(self, index_symbol, *, date):
|
|
return ["600000.XSHG", "000001.XSHE"]
|
|
|
|
def get_extras(
|
|
self,
|
|
info,
|
|
security_list,
|
|
*,
|
|
start_date,
|
|
end_date,
|
|
df,
|
|
):
|
|
self.last_get_extras_kwargs = {
|
|
"info": info,
|
|
"security_list": list(security_list),
|
|
"start_date": start_date,
|
|
"end_date": end_date,
|
|
"df": df,
|
|
}
|
|
return FakeFrame(
|
|
[
|
|
(
|
|
start_date,
|
|
{symbol: False for symbol in security_list},
|
|
),
|
|
(
|
|
end_date,
|
|
{symbol: False for symbol in security_list},
|
|
),
|
|
]
|
|
)
|
|
|
|
def get_bars(self, symbol, **kwargs):
|
|
self.last_get_bars_kwargs = dict(kwargs)
|
|
offset = 0 if symbol.startswith("600000") else 10
|
|
return FakeFrame(
|
|
[
|
|
(
|
|
0,
|
|
{
|
|
"date": kwargs["start_dt"],
|
|
"open": 10 + offset,
|
|
"high": 11 + offset,
|
|
"low": 9 + offset,
|
|
"close": 10.5 + offset,
|
|
"volume": 1000,
|
|
"money": 10_500 + offset * 1000,
|
|
"factor": 1.0,
|
|
"paused": False,
|
|
"high_limit": 11.5 + offset,
|
|
"low_limit": 8.5 + offset,
|
|
"pre_close": 10 + offset,
|
|
},
|
|
),
|
|
(
|
|
1,
|
|
{
|
|
"date": kwargs["end_dt"],
|
|
"open": 10.5 + offset,
|
|
"high": 11.5 + offset,
|
|
"low": 10 + offset,
|
|
"close": 11 + offset,
|
|
"volume": 1200,
|
|
"money": 13_200 + offset * 1200,
|
|
"factor": 1.0,
|
|
"paused": False,
|
|
"high_limit": 12 + offset,
|
|
"low_limit": 9 + offset,
|
|
"pre_close": 10.5 + offset,
|
|
},
|
|
),
|
|
]
|
|
)
|
|
|
|
|
|
class DataJQDataTests(unittest.TestCase):
|
|
def test_authentication_never_accepts_failed_login(self):
|
|
sdk = FakeJQData()
|
|
with self.assertRaisesRegex(JQDataUnavailableError, "not accepted"):
|
|
authenticate(
|
|
sdk,
|
|
username="user",
|
|
password="wrong",
|
|
interactive=False,
|
|
)
|
|
authenticate(
|
|
sdk,
|
|
username="user",
|
|
password="password",
|
|
interactive=False,
|
|
)
|
|
self.assertTrue(sdk.authenticated)
|
|
|
|
def test_fake_jqdata_snapshot_is_canonical_and_verifiable(self):
|
|
sdk = FakeJQData()
|
|
authenticate(
|
|
sdk,
|
|
username="user",
|
|
password="password",
|
|
interactive=False,
|
|
)
|
|
with tempfile.TemporaryDirectory() as directory:
|
|
paths = fetch_index_daily_snapshot(
|
|
sdk,
|
|
index_symbol="000905.XSHG",
|
|
start_date=date(2024, 1, 2),
|
|
end_date=date(2024, 1, 3),
|
|
output_dir=directory,
|
|
retrieved_at=datetime(2024, 1, 4, tzinfo=timezone.utc),
|
|
)
|
|
verified = verify_data_snapshot(paths["manifest"])
|
|
self.assertTrue(verified["ok"])
|
|
self.assertEqual(verified["bar_count"], 4)
|
|
self.assertEqual(verified["membership_count"], 4)
|
|
self.assertEqual(
|
|
sdk.last_get_bars_kwargs["end_dt"].date(),
|
|
date(2024, 1, 3),
|
|
)
|
|
self.assertEqual(
|
|
sdk.last_get_bars_kwargs["end_dt"].time(),
|
|
datetime.strptime("23:59:59", "%H:%M:%S").time(),
|
|
)
|
|
self.assertIs(
|
|
sdk.last_get_bars_kwargs["include_now"],
|
|
True,
|
|
)
|
|
self.assertEqual(
|
|
sdk.last_get_extras_kwargs["info"],
|
|
"is_st",
|
|
)
|
|
manifest = json.loads(
|
|
Path(paths["manifest"]).read_text(encoding="utf-8")
|
|
)
|
|
self.assertEqual(
|
|
manifest["next_trading_session"],
|
|
"2024-01-04",
|
|
)
|
|
self.assertEqual(
|
|
manifest["query"]["trading_sessions"],
|
|
["2024-01-02", "2024-01-03", "2024-01-04"],
|
|
)
|
|
self.assertEqual(sdk.full_calendar_queries, 1)
|
|
content = Path(paths["manifest"]).read_text(encoding="utf-8")
|
|
self.assertNotIn("password", content)
|
|
self.assertNotIn("user", content)
|
|
bars = (
|
|
Path(directory) / "bars.jsonl"
|
|
).read_text(encoding="utf-8")
|
|
self.assertIn('"adjustment_factor":1.0', bars)
|
|
self.assertIn('"limit_up":', bars)
|
|
|
|
def test_next_session_comes_from_provider_calendar_across_holiday(self):
|
|
class NationalDayCalendarJQData(FakeJQData):
|
|
def get_all_trade_days(self):
|
|
self.full_calendar_queries += 1
|
|
return [
|
|
date(2024, 9, 27),
|
|
date(2024, 9, 30),
|
|
date(2024, 10, 8),
|
|
date(2024, 10, 9),
|
|
]
|
|
|
|
sdk = NationalDayCalendarJQData()
|
|
with tempfile.TemporaryDirectory() as directory:
|
|
paths = fetch_index_daily_snapshot(
|
|
sdk,
|
|
index_symbol="000905.XSHG",
|
|
start_date=date(2024, 9, 27),
|
|
end_date=date(2024, 9, 30),
|
|
output_dir=directory,
|
|
retrieved_at=datetime(2024, 10, 9, tzinfo=timezone.utc),
|
|
)
|
|
manifest = json.loads(
|
|
Path(paths["manifest"]).read_text(encoding="utf-8")
|
|
)
|
|
self.assertEqual(
|
|
manifest["next_trading_session"],
|
|
"2024-10-08",
|
|
)
|
|
self.assertEqual(
|
|
manifest["query"]["trading_sessions"][-1],
|
|
"2024-10-08",
|
|
)
|
|
|
|
def test_invalid_next_session_calendar_response_fails_closed(self):
|
|
class BadCalendarJQData(FakeJQData):
|
|
def get_all_trade_days(self):
|
|
return [date(2024, 1, 2), date(2024, 1, 3)]
|
|
|
|
with tempfile.TemporaryDirectory() as directory:
|
|
with self.assertRaisesRegex(
|
|
JQDataUnavailableError,
|
|
"unique next trading session",
|
|
):
|
|
fetch_index_daily_snapshot(
|
|
BadCalendarJQData(),
|
|
index_symbol="000905.XSHG",
|
|
start_date=date(2024, 1, 2),
|
|
end_date=date(2024, 1, 3),
|
|
output_dir=directory,
|
|
retrieved_at=datetime(
|
|
2024,
|
|
1,
|
|
4,
|
|
tzinfo=timezone.utc,
|
|
),
|
|
)
|
|
|
|
def test_inconsistent_range_and_full_calendars_fail_closed(self):
|
|
class MissingMiddleSessionJQData(FakeJQData):
|
|
def get_all_trade_days(self):
|
|
return [
|
|
date(2024, 1, 2),
|
|
date(2024, 1, 3),
|
|
date(2024, 1, 4),
|
|
]
|
|
|
|
def get_trade_days(self, *, start_date, end_date):
|
|
return [start_date]
|
|
|
|
with tempfile.TemporaryDirectory() as directory:
|
|
with self.assertRaisesRegex(
|
|
JQDataUnavailableError,
|
|
"inconsistent sessions",
|
|
):
|
|
fetch_index_daily_snapshot(
|
|
MissingMiddleSessionJQData(),
|
|
index_symbol="000905.XSHG",
|
|
start_date=date(2024, 1, 2),
|
|
end_date=date(2024, 1, 3),
|
|
output_dir=directory,
|
|
retrieved_at=datetime(
|
|
2024,
|
|
1,
|
|
4,
|
|
tzinfo=timezone.utc,
|
|
),
|
|
)
|
|
|
|
def test_same_shanghai_day_is_rejected_before_query(self):
|
|
sdk = FakeJQData()
|
|
with tempfile.TemporaryDirectory() as directory:
|
|
with self.assertRaisesRegex(
|
|
JQDataUnavailableError,
|
|
"24:00 finalization",
|
|
):
|
|
fetch_index_daily_snapshot(
|
|
sdk,
|
|
index_symbol="000905.XSHG",
|
|
start_date=date(2024, 1, 2),
|
|
end_date=date(2024, 1, 3),
|
|
output_dir=directory,
|
|
retrieved_at=datetime(
|
|
2024,
|
|
1,
|
|
3,
|
|
6,
|
|
tzinfo=timezone.utc,
|
|
),
|
|
)
|
|
self.assertIsNone(sdk.last_get_bars_kwargs)
|
|
self.assertIsNone(sdk.last_get_extras_kwargs)
|
|
|
|
def test_missing_pit_st_status_fails_closed(self):
|
|
class MissingStatusJQData(FakeJQData):
|
|
def get_extras(self, *args, **kwargs):
|
|
frame = super().get_extras(*args, **kwargs)
|
|
frame.rows[-1][1].pop("600000.XSHG")
|
|
return frame
|
|
|
|
with tempfile.TemporaryDirectory() as directory:
|
|
with self.assertRaisesRegex(
|
|
JQDataUnavailableError,
|
|
"omitted required field 600000.XSHG",
|
|
):
|
|
fetch_index_daily_snapshot(
|
|
MissingStatusJQData(),
|
|
index_symbol="000905.XSHG",
|
|
start_date=date(2024, 1, 2),
|
|
end_date=date(2024, 1, 3),
|
|
output_dir=directory,
|
|
retrieved_at=datetime(
|
|
2024,
|
|
1,
|
|
4,
|
|
tzinfo=timezone.utc,
|
|
),
|
|
)
|
|
|
|
def test_missing_required_provider_rule_field_fails_closed(self):
|
|
class MissingLimitJQData(FakeJQData):
|
|
def get_bars(self, symbol, **kwargs):
|
|
frame = super().get_bars(symbol, **kwargs)
|
|
del frame.rows[0][1]["high_limit"]
|
|
return frame
|
|
|
|
with tempfile.TemporaryDirectory() as directory:
|
|
with self.assertRaisesRegex(
|
|
JQDataUnavailableError,
|
|
"omitted required field high_limit",
|
|
):
|
|
fetch_index_daily_snapshot(
|
|
MissingLimitJQData(),
|
|
index_symbol="000905.XSHG",
|
|
start_date=date(2024, 1, 2),
|
|
end_date=date(2024, 1, 3),
|
|
output_dir=directory,
|
|
retrieved_at=datetime(
|
|
2024,
|
|
1,
|
|
4,
|
|
tzinfo=timezone.utc,
|
|
),
|
|
)
|
|
|
|
def test_membership_without_same_session_bar_fails_quality_gate(self):
|
|
class MissingEndBarJQData(FakeJQData):
|
|
def get_bars(self, symbol, **kwargs):
|
|
frame = super().get_bars(symbol, **kwargs)
|
|
frame.rows.pop()
|
|
return frame
|
|
|
|
with tempfile.TemporaryDirectory() as directory:
|
|
with self.assertRaisesRegex(
|
|
DataQualityError,
|
|
"membership has no same-session canonical bar",
|
|
):
|
|
fetch_index_daily_snapshot(
|
|
MissingEndBarJQData(),
|
|
index_symbol="000905.XSHG",
|
|
start_date=date(2024, 1, 2),
|
|
end_date=date(2024, 1, 3),
|
|
output_dir=directory,
|
|
retrieved_at=datetime(
|
|
2024,
|
|
1,
|
|
4,
|
|
tzinfo=timezone.utc,
|
|
),
|
|
)
|
|
|
|
def test_bad_ohlc_fails_quality_gate(self):
|
|
with self.assertRaisesRegex(DataQualityError, "open must lie"):
|
|
CanonicalBarRecord(
|
|
trading_date=date(2024, 1, 2),
|
|
symbol="600000.SH",
|
|
open=12,
|
|
high=11,
|
|
low=9,
|
|
close=10,
|
|
volume=100,
|
|
money=1000,
|
|
paused=False,
|
|
source="fixture",
|
|
)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
unittest.main()
|