You cannot select more than 25 topics
Topics must start with a letter or number, can include dashes ('-') and can be up to 35 characters long.
500 lines
14 KiB
Python
500 lines
14 KiB
Python
from __future__ import annotations
|
|
|
|
from concurrent.futures import ThreadPoolExecutor
|
|
import multiprocessing
|
|
from pathlib import Path
|
|
import threading
|
|
import time
|
|
|
|
import pyarrow as pa
|
|
import pyarrow.parquet as pq
|
|
|
|
from libs.intraday.cache import DailyBarCache, IntradayCache, LayeredDailyBarCache
|
|
|
|
|
|
def _daily_cache_put_worker(
|
|
cache_dir: str,
|
|
start_date: str,
|
|
end_date: str,
|
|
bars: list[dict[str, float | str]],
|
|
start_event,
|
|
read_delay_s: float,
|
|
) -> None:
|
|
cache = DailyBarCache(cache_dir)
|
|
target_path = str(Path(cache_dir) / "AAPL.parquet")
|
|
original_read_table = pq.read_table
|
|
|
|
def delayed_read_table(path, *args, **kwargs):
|
|
table = original_read_table(path, *args, **kwargs)
|
|
if str(path) == target_path:
|
|
time.sleep(read_delay_s)
|
|
return table
|
|
|
|
pq.read_table = delayed_read_table
|
|
start_event.wait(timeout=5.0)
|
|
cache.put("AAPL", start_date, end_date, bars)
|
|
|
|
|
|
def test_legacy_cache_file_is_treated_as_miss_and_removed(tmp_path) -> None:
|
|
cache = IntradayCache(str(tmp_path))
|
|
cache_path = tmp_path / "AAPL" / "2026-01-05.parquet"
|
|
cache_path.parent.mkdir(parents=True, exist_ok=True)
|
|
|
|
legacy = pa.table({
|
|
"timestamp": ["2026-01-05T13:30:00Z"],
|
|
"open": [100.0],
|
|
"high": [101.0],
|
|
"low": [99.5],
|
|
"close": [100.5],
|
|
"volume": [1000.0],
|
|
"vwap": [100.4],
|
|
})
|
|
pq.write_table(legacy, str(cache_path))
|
|
|
|
assert cache.has("AAPL", "2026-01-05") is False
|
|
assert cache.get("AAPL", "2026-01-05") is None
|
|
assert cache_path.exists() is False
|
|
|
|
|
|
def test_cache_put_writes_metadata_valid_file(tmp_path) -> None:
|
|
cache = IntradayCache(str(tmp_path))
|
|
bars = []
|
|
for i in range(10):
|
|
bars.append({
|
|
"timestamp": f"2026-01-05T13:{30 + i:02d}:00Z",
|
|
"open": 100.0,
|
|
"high": 101.0,
|
|
"low": 99.5,
|
|
"close": 100.5,
|
|
"volume": 1000.0,
|
|
"vwap": 100.4,
|
|
})
|
|
|
|
cache.put("AAPL", "2026-01-05", bars)
|
|
|
|
assert cache.has("AAPL", "2026-01-05") is True
|
|
assert cache.get("AAPL", "2026-01-05") == bars
|
|
|
|
|
|
def test_intraday_cache_rejects_too_few_rows(tmp_path) -> None:
|
|
cache = IntradayCache(str(tmp_path))
|
|
cache_path = tmp_path / "AAPL" / "2026-01-05.parquet"
|
|
cache_path.parent.mkdir(parents=True, exist_ok=True)
|
|
|
|
tiny = pa.table({
|
|
"timestamp": [f"2026-01-05T13:{30+i:02d}:00Z" for i in range(5)],
|
|
"open": [100.0] * 5,
|
|
"high": [101.0] * 5,
|
|
"low": [99.5] * 5,
|
|
"close": [100.5] * 5,
|
|
"volume": [1000.0] * 5,
|
|
"vwap": [100.4] * 5,
|
|
}, schema=pa.schema([
|
|
pa.field("timestamp", pa.string()),
|
|
pa.field("open", pa.float64()),
|
|
pa.field("high", pa.float64()),
|
|
pa.field("low", pa.float64()),
|
|
pa.field("close", pa.float64()),
|
|
pa.field("volume", pa.float64()),
|
|
pa.field("vwap", pa.float64()),
|
|
]).with_metadata({
|
|
b"intraday_cache_version": b"3",
|
|
b"intraday_cache_source": b"api_v1_alpaca_intraday",
|
|
b"intraday_cache_interval": b"5min",
|
|
}))
|
|
pq.write_table(tiny, str(cache_path))
|
|
|
|
assert cache.has("AAPL", "2026-01-05") is False
|
|
assert cache_path.exists() is False
|
|
|
|
|
|
def test_intraday_cache_accepts_existing_positive_file_without_kind_metadata(tmp_path) -> None:
|
|
cache = IntradayCache(str(tmp_path))
|
|
cache_path = tmp_path / "AAPL" / "2026-01-05.parquet"
|
|
cache_path.parent.mkdir(parents=True, exist_ok=True)
|
|
|
|
rows = 12
|
|
table = pa.table({
|
|
"timestamp": [f"2026-01-05T13:{30+i:02d}:00Z" for i in range(rows)],
|
|
"open": [100.0] * rows,
|
|
"high": [101.0] * rows,
|
|
"low": [99.5] * rows,
|
|
"close": [100.5] * rows,
|
|
"volume": [1000.0] * rows,
|
|
"vwap": [100.4] * rows,
|
|
}, schema=pa.schema([
|
|
pa.field("timestamp", pa.string()),
|
|
pa.field("open", pa.float64()),
|
|
pa.field("high", pa.float64()),
|
|
pa.field("low", pa.float64()),
|
|
pa.field("close", pa.float64()),
|
|
pa.field("volume", pa.float64()),
|
|
pa.field("vwap", pa.float64()),
|
|
]).with_metadata({
|
|
b"intraday_cache_version": b"3",
|
|
b"intraday_cache_source": b"api_v1_alpaca_intraday",
|
|
b"intraday_cache_interval": b"5min",
|
|
}))
|
|
pq.write_table(table, str(cache_path))
|
|
|
|
assert cache.has("AAPL", "2026-01-05") is True
|
|
assert len(cache.get("AAPL", "2026-01-05") or []) == rows
|
|
|
|
|
|
def test_intraday_cache_negative_entry_suppresses_rereads(tmp_path) -> None:
|
|
cache = IntradayCache(str(tmp_path))
|
|
|
|
cache.put_negative("AAPL", "2026-01-05", reason="sparse")
|
|
|
|
assert cache.has("AAPL", "2026-01-05") is True
|
|
assert cache.get("AAPL", "2026-01-05") == []
|
|
|
|
|
|
def test_daily_cache_put_and_get_uses_requested_coverage_range(tmp_path) -> None:
|
|
cache = DailyBarCache(str(tmp_path))
|
|
bars = [
|
|
{
|
|
"date": "2026-01-02",
|
|
"open": 100.0,
|
|
"high": 101.0,
|
|
"low": 99.0,
|
|
"close": 100.5,
|
|
"volume": 1000.0,
|
|
},
|
|
{
|
|
"date": "2026-01-05",
|
|
"open": 101.0,
|
|
"high": 102.0,
|
|
"low": 100.0,
|
|
"close": 101.5,
|
|
"volume": 1200.0,
|
|
},
|
|
]
|
|
|
|
cache.put("AAPL", "2026-01-01", "2026-01-10", bars)
|
|
|
|
assert cache.get("AAPL", "2026-01-01", "2026-01-10") == bars
|
|
assert cache.get("AAPL", "2026-01-02", "2026-01-05") == bars
|
|
assert cache.get("AAPL", "2025-12-31", "2026-01-10") is None
|
|
|
|
|
|
def test_daily_cache_treats_egregiously_partial_warmup_file_as_miss_without_deleting(tmp_path) -> None:
|
|
cache = DailyBarCache(str(tmp_path))
|
|
cache_path = tmp_path / "BLD.parquet"
|
|
rows = []
|
|
for i in range(11):
|
|
rows.append(
|
|
{
|
|
"date": f"2026-04-{6 + i:02d}",
|
|
"open": 100.0 + i,
|
|
"high": 101.0 + i,
|
|
"low": 99.0 + i,
|
|
"close": 100.5 + i,
|
|
"volume": 1_000.0 + i,
|
|
}
|
|
)
|
|
|
|
schema = pa.schema([
|
|
pa.field("date", pa.string()),
|
|
pa.field("open", pa.float64()),
|
|
pa.field("high", pa.float64()),
|
|
pa.field("low", pa.float64()),
|
|
pa.field("close", pa.float64()),
|
|
pa.field("volume", pa.float64()),
|
|
]).with_metadata({
|
|
b"daily_cache_version": b"1",
|
|
b"daily_cache_source": b"api_v1_price_data",
|
|
b"daily_cache_interval": b"1d",
|
|
b"daily_cache_coverage_start": b"2026-01-20",
|
|
b"daily_cache_coverage_end": b"2026-04-20",
|
|
})
|
|
pq.write_table(pa.Table.from_pylist(rows, schema=schema), str(cache_path))
|
|
|
|
assert cache.get("BLD", "2026-01-20", "2026-04-20") is None
|
|
assert cache_path.exists() is True
|
|
cached_rows, tail_start = cache.get_with_tail("BLD", "2026-01-20", "2026-04-20")
|
|
assert cached_rows is None
|
|
assert tail_start is None
|
|
|
|
|
|
def test_daily_cache_put_tightens_sparse_fallback_coverage_to_actual_rows(tmp_path) -> None:
|
|
cache = DailyBarCache(str(tmp_path))
|
|
rows = [
|
|
{
|
|
"date": f"2026-04-{6 + i:02d}",
|
|
"open": 100.0 + i,
|
|
"high": 101.0 + i,
|
|
"low": 99.0 + i,
|
|
"close": 100.5 + i,
|
|
"volume": 1_000.0 + i,
|
|
}
|
|
for i in range(11)
|
|
]
|
|
|
|
cache.put("BLD", "2026-01-20", "2026-04-20", rows)
|
|
|
|
cache_path = tmp_path / "BLD.parquet"
|
|
metadata = pq.read_metadata(str(cache_path)).schema.to_arrow_schema().metadata or {}
|
|
assert metadata[b"daily_cache_coverage_start"] == b"2026-04-06"
|
|
assert metadata[b"daily_cache_coverage_end"] == b"2026-04-16"
|
|
cached_rows, tail_start = cache.get_with_tail("BLD", "2026-01-20", "2026-04-20")
|
|
assert cached_rows is None
|
|
assert tail_start is None
|
|
cached_rows, tail_start = cache.get_with_tail("BLD", "2026-04-06", "2026-04-16")
|
|
assert cached_rows == rows
|
|
assert tail_start is None
|
|
|
|
|
|
def test_daily_cache_accepts_late_start_when_history_is_sufficient(tmp_path) -> None:
|
|
cache = DailyBarCache(str(tmp_path))
|
|
rows = [
|
|
{
|
|
"date": f"2026-04-{1 + i:02d}",
|
|
"open": 100.0 + i,
|
|
"high": 101.0 + i,
|
|
"low": 99.0 + i,
|
|
"close": 100.5 + i,
|
|
"volume": 1_000.0 + i,
|
|
}
|
|
for i in range(20)
|
|
]
|
|
|
|
cache.put("IPO", "2026-01-20", "2026-04-20", rows)
|
|
|
|
assert cache.get("IPO", "2026-01-20", "2026-04-20") == rows
|
|
cached_rows, tail_start = cache.get_with_tail("IPO", "2026-01-20", "2026-04-20")
|
|
assert cached_rows == rows
|
|
assert tail_start is None
|
|
|
|
|
|
def test_layered_daily_cache_reads_snapshot_and_writes_overlay(tmp_path) -> None:
|
|
snapshot_cache = DailyBarCache(str(tmp_path / "snapshot"))
|
|
overlay_cache = DailyBarCache(str(tmp_path / "overlay"))
|
|
layered = LayeredDailyBarCache(snapshot_cache, overlay_cache)
|
|
rows = [
|
|
{
|
|
"date": "2026-01-02",
|
|
"open": 100.0,
|
|
"high": 101.0,
|
|
"low": 99.0,
|
|
"close": 100.5,
|
|
"volume": 1000.0,
|
|
}
|
|
]
|
|
|
|
snapshot_cache.put("AAPL", "2026-01-01", "2026-01-10", rows)
|
|
|
|
assert layered.get("AAPL", "2026-01-01", "2026-01-10") == rows
|
|
|
|
extended_rows = rows + [
|
|
{
|
|
"date": "2026-01-11",
|
|
"open": 101.0,
|
|
"high": 102.0,
|
|
"low": 100.0,
|
|
"close": 101.5,
|
|
"volume": 1100.0,
|
|
}
|
|
]
|
|
layered.put("AAPL", "2026-01-01", "2026-01-11", extended_rows)
|
|
|
|
assert snapshot_cache.get("AAPL", "2026-01-01", "2026-01-10") == rows
|
|
assert snapshot_cache.get("AAPL", "2026-01-01", "2026-01-11") is None
|
|
assert overlay_cache.get("AAPL", "2026-01-01", "2026-01-11") == extended_rows
|
|
|
|
|
|
def test_daily_cache_merges_concurrent_puts_without_losing_coverage(tmp_path, monkeypatch) -> None:
|
|
cache = DailyBarCache(str(tmp_path))
|
|
cache.put(
|
|
"AAPL",
|
|
"2026-01-01",
|
|
"2026-01-05",
|
|
[
|
|
{
|
|
"date": "2026-01-02",
|
|
"open": 100.0,
|
|
"high": 101.0,
|
|
"low": 99.0,
|
|
"close": 100.5,
|
|
"volume": 1000.0,
|
|
}
|
|
],
|
|
)
|
|
|
|
target_path = str((tmp_path / "AAPL.parquet"))
|
|
original_read_table = pq.read_table
|
|
read_count = 0
|
|
read_count_lock = threading.Lock()
|
|
allow_reads_to_continue = threading.Event()
|
|
|
|
def delayed_read_table(path, *args, **kwargs):
|
|
nonlocal read_count
|
|
table = original_read_table(path, *args, **kwargs)
|
|
if str(path) == target_path:
|
|
with read_count_lock:
|
|
read_count += 1
|
|
if read_count >= 2:
|
|
allow_reads_to_continue.set()
|
|
allow_reads_to_continue.wait(timeout=0.2)
|
|
return table
|
|
|
|
monkeypatch.setattr(pq, "read_table", delayed_read_table)
|
|
|
|
first_extension = [
|
|
{
|
|
"date": "2026-01-06",
|
|
"open": 101.0,
|
|
"high": 102.0,
|
|
"low": 100.0,
|
|
"close": 101.5,
|
|
"volume": 1100.0,
|
|
}
|
|
]
|
|
second_extension = [
|
|
{
|
|
"date": "2026-01-07",
|
|
"open": 102.0,
|
|
"high": 103.0,
|
|
"low": 101.0,
|
|
"close": 102.5,
|
|
"volume": 1200.0,
|
|
}
|
|
]
|
|
|
|
with ThreadPoolExecutor(max_workers=2) as pool:
|
|
futures = [
|
|
pool.submit(cache.put, "AAPL", "2026-01-01", "2026-01-06", first_extension),
|
|
pool.submit(cache.put, "AAPL", "2026-01-01", "2026-01-07", second_extension),
|
|
]
|
|
for future in futures:
|
|
future.result()
|
|
|
|
assert cache.get("AAPL", "2026-01-01", "2026-01-07") == [
|
|
{
|
|
"date": "2026-01-02",
|
|
"open": 100.0,
|
|
"high": 101.0,
|
|
"low": 99.0,
|
|
"close": 100.5,
|
|
"volume": 1000.0,
|
|
},
|
|
{
|
|
"date": "2026-01-06",
|
|
"open": 101.0,
|
|
"high": 102.0,
|
|
"low": 100.0,
|
|
"close": 101.5,
|
|
"volume": 1100.0,
|
|
},
|
|
{
|
|
"date": "2026-01-07",
|
|
"open": 102.0,
|
|
"high": 103.0,
|
|
"low": 101.0,
|
|
"close": 102.5,
|
|
"volume": 1200.0,
|
|
},
|
|
]
|
|
|
|
|
|
def test_daily_cache_merges_cross_process_puts_without_losing_coverage(tmp_path) -> None:
|
|
cache = DailyBarCache(str(tmp_path))
|
|
cache.put(
|
|
"AAPL",
|
|
"2026-01-01",
|
|
"2026-01-05",
|
|
[
|
|
{
|
|
"date": "2026-01-02",
|
|
"open": 100.0,
|
|
"high": 101.0,
|
|
"low": 99.0,
|
|
"close": 100.5,
|
|
"volume": 1000.0,
|
|
}
|
|
],
|
|
)
|
|
|
|
first_extension = [
|
|
{
|
|
"date": "2026-01-06",
|
|
"open": 101.0,
|
|
"high": 102.0,
|
|
"low": 100.0,
|
|
"close": 101.5,
|
|
"volume": 1100.0,
|
|
}
|
|
]
|
|
second_extension = [
|
|
{
|
|
"date": "2026-01-07",
|
|
"open": 102.0,
|
|
"high": 103.0,
|
|
"low": 101.0,
|
|
"close": 102.5,
|
|
"volume": 1200.0,
|
|
}
|
|
]
|
|
|
|
ctx = multiprocessing.get_context("spawn")
|
|
start_event = ctx.Event()
|
|
processes = [
|
|
ctx.Process(
|
|
target=_daily_cache_put_worker,
|
|
args=(
|
|
str(tmp_path),
|
|
"2026-01-01",
|
|
"2026-01-06",
|
|
first_extension,
|
|
start_event,
|
|
0.3,
|
|
),
|
|
),
|
|
ctx.Process(
|
|
target=_daily_cache_put_worker,
|
|
args=(
|
|
str(tmp_path),
|
|
"2026-01-01",
|
|
"2026-01-07",
|
|
second_extension,
|
|
start_event,
|
|
0.3,
|
|
),
|
|
),
|
|
]
|
|
|
|
for process in processes:
|
|
process.start()
|
|
start_event.set()
|
|
|
|
for process in processes:
|
|
process.join(timeout=10.0)
|
|
assert process.exitcode == 0
|
|
|
|
assert cache.get("AAPL", "2026-01-01", "2026-01-07") == [
|
|
{
|
|
"date": "2026-01-02",
|
|
"open": 100.0,
|
|
"high": 101.0,
|
|
"low": 99.0,
|
|
"close": 100.5,
|
|
"volume": 1000.0,
|
|
},
|
|
{
|
|
"date": "2026-01-06",
|
|
"open": 101.0,
|
|
"high": 102.0,
|
|
"low": 100.0,
|
|
"close": 101.5,
|
|
"volume": 1100.0,
|
|
},
|
|
{
|
|
"date": "2026-01-07",
|
|
"open": 102.0,
|
|
"high": 103.0,
|
|
"low": 101.0,
|
|
"close": 102.5,
|
|
"volume": 1200.0,
|
|
},
|
|
]
|