Repository Public
1from __future__ import annotations23import shutil4import socket5import subprocess67import numpy as np8import pandas as pd9import pytest1011from conftest import make_series12from spq.data import DataError, align, bars_per_year, interval_seconds, load_candles, load_universe, normalise13from spq.data import fetch as fetch_mod1415AS_OF = "2030-01-01"161718def test_interval_and_annualisation():19 assert interval_seconds("1m") == 60 and interval_seconds("4h") == 14400 and interval_seconds("1d") == 8640020 assert bars_per_year("1h") == 8760 and bars_per_year("1d") == 365 and bars_per_year("15m") == 3504021 with pytest.raises(DataError):22 interval_seconds("1x")232425def test_normalise_sorts_and_keeps_closed_bars_only():26 df = make_series(10).reset_index()27 shuffled = df.sample(frac=1.0, random_state=1)28 as_of = df["ts"].iloc[-1] + pd.Timedelta(minutes=30) # the last bar is still open at as_of29 out, q = normalise(shuffled, symbol="X", interval="1h", source="t", as_of=as_of)30 assert out.index.is_monotonic_increasing and str(out.index.tz) == "UTC"31 assert len(out) == 9 and q.unclosed_dropped == 132 assert out.index[-1] == df["ts"].iloc[-2]33 assert list(out.columns) == ["open", "high", "low", "close", "volume"]343536def test_bar_closing_exactly_at_as_of_is_closed():37 df = make_series(5).reset_index()38 out, _ = normalise(df, symbol="X", interval="1h", source="t", as_of=df["ts"].iloc[-1] + pd.Timedelta(hours=1))39 assert len(out) == 5404142def test_duplicates_raise_and_identical_can_be_dropped():43 df = make_series(6).reset_index()44 dup = pd.concat([df, df.iloc[[2]]])45 with pytest.raises(DataError, match="share a timestamp"):46 normalise(dup, symbol="X", interval="1h", source="t", as_of=AS_OF)47 out, q = normalise(dup, symbol="X", interval="1h", source="t", as_of=AS_OF, on_duplicate="drop_identical")48 assert len(out) == 6 and q.duplicates_dropped == 149 conflict = df.iloc[[2]].copy()50 conflict["close"] *= 1.0151 conflict["high"] *= 1.0252 with pytest.raises(DataError, match="conflicting"):53 normalise(pd.concat([df, conflict]), symbol="X", interval="1h", source="t", as_of=AS_OF,54 on_duplicate="drop_identical")555657def test_gaps_are_reported_not_filled():58 df = make_series(20).reset_index().drop(index=[5, 6, 7, 12])59 out, q = normalise(df, symbol="X", interval="1h", source="t", as_of=AS_OF)60 assert len(out) == 16 and q.missing_bars == 4 and q.max_gap_bars == 3 and q.expected_bars == 2061 assert [g[2] for g in q.gaps] == [3, 1]62 assert q.gaps[0][0] == "2024-01-01T05:00:00+00:00" and q.gaps[0][1] == "2024-01-01T07:00:00+00:00"63 assert not q.ok() and q.ok(max_missing_fraction=0.5)64 with pytest.raises(DataError, match="missing bars"):65 normalise(df, symbol="X", interval="1h", source="t", as_of=AS_OF, on_gap="raise")666768def test_ohlc_violations_and_off_grid_fail_quality():69 df = make_series(10).reset_index()70 df.loc[3, "high"] = df.loc[3, "low"] * 0.571 df.loc[4, "close"] = np.nan72 _, q = normalise(df, symbol="X", interval="1h", source="t", as_of=AS_OF)73 assert q.ohlc_violations == 2 and not q.ok()74 df = make_series(10).reset_index()75 df["ts"] = df["ts"] + pd.Timedelta(minutes=7)76 _, q = normalise(df, symbol="X", interval="1h", source="t", as_of=AS_OF)77 assert q.off_grid == 10 and not q.ok()787980def test_epoch_and_naive_timestamps_are_utc():81 df = make_series(5).reset_index()82 expected = df.set_index("ts").index83 secs = df.assign(ts=[int(t.timestamp()) for t in df["ts"]])84 out, _ = normalise(secs, symbol="X", interval="1h", source="t", as_of=AS_OF)85 assert (out.index == expected).all() and str(out.index.tz) == "UTC"86 millis = df.assign(ts=[int(t.timestamp()) * 1000 for t in df["ts"]])87 out, _ = normalise(millis, symbol="X", interval="1h", source="t", as_of=AS_OF)88 assert (out.index == expected).all()89 naive = df.assign(ts=df["ts"].dt.tz_localize(None))90 out, _ = normalise(naive, symbol="X", interval="1h", source="t", as_of=AS_OF)91 assert (out.index == expected).all() and str(out.index.tz) == "UTC"929394@pytest.mark.parametrize("ext", ["parquet", "csv"])95def test_files_roundtrip_with_range(tmp_path, ext):96 df = make_series(48)97 path = tmp_path / "coinbase" / f"ETH_1h.{ext}"98 path.parent.mkdir(parents=True)99 df.reset_index().to_parquet(path, index=False) if ext == "parquet" else df.reset_index().to_csv(path, index=False)100 out, q = load_candles("eth", "1h", "2024-01-01T10:00", "2024-01-01T20:00", root=tmp_path, as_of=AS_OF)101 assert len(out) == 10 and out.index[0] == pd.Timestamp("2024-01-01T10:00", tz="UTC")102 assert out.index[-1] == pd.Timestamp("2024-01-01T19:00", tz="UTC") # end is exclusive103 np.testing.assert_allclose(out["close"].to_numpy(), df["close"].iloc[10:20].to_numpy())104 assert q.source == "coinbase" and q.ok()105106107def test_missing_origin_and_file(tmp_path, monkeypatch):108 monkeypatch.delenv("SPQ_DATA_DIR", raising=False)109 monkeypatch.delenv("DATABASE_URL", raising=False)110 with pytest.raises(DataError, match="no data origin"):111 load_candles("ETH", "1h")112 with pytest.raises(DataError, match="no candle file"):113 load_candles("ETH", "1h", root=tmp_path)114115116def test_universe_is_aligned_on_common_bars(tmp_path):117 a, b = make_series(30, seed=1), make_series(30, seed=2).drop(index=make_series(30).index[[4, 9]])118 for name, df in (("AAA", a), ("BBB", b)):119 p = tmp_path / "coinbase" / f"{name}_1h.parquet"120 p.parent.mkdir(exist_ok=True)121 df.reset_index().to_parquet(p, index=False)122 data, quality = load_universe(["AAA", "BBB"], "1h", root=tmp_path, as_of=AS_OF)123 assert len(data["AAA"]) == len(data["BBB"]) == 28 and data["AAA"].index.equals(data["BBB"].index)124 assert quality["BBB"].missing_bars == 2 and quality["AAA"].missing_bars == 0125 assert align({}) == {}126127128# ------------------------------------------------------------------------------------------------- downloader129130def test_coinbase_parser_orders_and_drops_open_bar(monkeypatch, tmp_path):131 now = pd.Timestamp.now(tz="UTC").floor("h")132 ts = [int((now - pd.Timedelta(hours=k)).timestamp()) for k in range(5)] # newest first, ts[0] still open133 rows = [[t, 9.0, 11.0, 10.0, 10.5, 3.0] for t in ts] # [time, low, high, open, close, volume]134 calls = []135136 def fake(url, params, **_):137 calls.append((url, params))138 return rows if len(calls) == 1 else []139140 monkeypatch.setattr(fetch_mod, "_get_json", fake)141 monkeypatch.setattr(fetch_mod.time, "sleep", lambda *_: None)142 path, q = fetch_mod.download("ETH", "1h", now - pd.Timedelta(hours=4), source="coinbase", out=tmp_path)143 assert "ETH-USD" in calls[0][0] and calls[0][1]["granularity"] == 3600144 df, _ = load_candles("ETH", "1h", root=tmp_path)145 assert len(df) == 4 and df.index.is_monotonic_increasing # the bar in progress is not stored146 assert df.index[-1] == now - pd.Timedelta(hours=1)147 assert df.iloc[0].to_dict() == {"open": 10.0, "high": 11.0, "low": 9.0, "close": 10.5, "volume": 3.0}148 # a second download merges without duplicating149 calls.clear()150 fetch_mod.download("ETH", "1h", now - pd.Timedelta(hours=4), source="coinbase", out=tmp_path)151 df2, q2 = load_candles("ETH", "1h", root=tmp_path)152 assert len(df2) == 4 and q2.duplicates_dropped == 0153154155def test_kraken_parser(monkeypatch, tmp_path):156 now = pd.Timestamp.now(tz="UTC").floor("h")157 ts = [int((now - pd.Timedelta(hours=k)).timestamp()) for k in range(3, -1, -1)] # oldest first, last one open158 rows = [[t, "10.0", "11.0", "9.0", "10.5", "10.2", "3.0", 7] for t in ts]159 seen = {}160161 def fake(url, params, **_):162 seen.update(params)163 return {"error": [], "result": {"XXBTZUSD": rows, "last": ts[-2]}}164165 monkeypatch.setattr(fetch_mod, "_get_json", fake)166 fetch_mod.download("BTC", "1h", now - pd.Timedelta(days=1), source="kraken", out=tmp_path)167 assert seen == {"pair": "XBTUSD", "interval": 60}168 df, _ = load_candles("BTC", "1h", root=tmp_path, source="kraken")169 assert len(df) == 3 and df["high"].iloc[0] == 11.0 and df["volume"].iloc[0] == 3.0170 monkeypatch.setattr(fetch_mod, "_get_json", lambda *a, **k: {"error": ["EQuery:Unknown asset pair"]})171 with pytest.raises(DataError, match="Unknown asset pair"):172 fetch_mod.fetch_kraken("NOPE", "1h")173174175@pytest.mark.network176def test_live_endpoints_answer():177 cb = fetch_mod.fetch_coinbase("ETH", "1h", pd.Timestamp.now(tz="UTC") - pd.Timedelta(hours=6))178 kr = fetch_mod.fetch_kraken("ETH", "1h")179 assert len(cb) >= 5 and len(kr) >= 700180181182# ------------------------------------------------------------------------------------------------- PostgreSQL183184@pytest.fixture(scope="module")185def pg_url(tmp_path_factory):186 initdb, pg_ctl = shutil.which("initdb"), shutil.which("pg_ctl")187 if not initdb or not pg_ctl:188 pytest.skip("PostgreSQL binaries (initdb, pg_ctl) not installed")189 base = tmp_path_factory.mktemp("pg")190 datadir = base / "d"191 subprocess.run([initdb, "-D", str(datadir), "-U", "spq", "--auth=trust", "-E", "UTF8"], check=True,192 capture_output=True)193 with socket.socket() as s:194 s.bind(("127.0.0.1", 0))195 port = s.getsockname()[1]196 opts = f"-p {port} -c listen_addresses=127.0.0.1 -c unix_socket_directories='' -c fsync=off"197 subprocess.run([pg_ctl, "-D", str(datadir), "-o", opts, "-l", str(base / "log"), "-w", "start"], check=True,198 capture_output=True)199 try:200 yield f"postgresql://spq@127.0.0.1:{port}/postgres"201 finally:202 subprocess.run([pg_ctl, "-D", str(datadir), "-m", "immediate", "stop"], capture_output=True)203204205def test_postgres_candles_table(pg_url):206 import psycopg207208 df = make_series(30)209 with psycopg.connect(pg_url, autocommit=True) as conn:210 conn.execute(211 "CREATE TABLE candles (source text NOT NULL, symbol text NOT NULL, interval text NOT NULL, "212 "ts timestamptz NOT NULL, open numeric NOT NULL, high numeric NOT NULL, low numeric NOT NULL, "213 "close numeric NOT NULL, volume numeric NOT NULL, PRIMARY KEY (source, symbol, interval, ts))")214 rows = [("coinbase", "ETH", "1h", ts.to_pydatetime(), *map(float, r)) for ts, r in215 zip(df.index, df.to_numpy())]216 rows += [("kraken", "ETH", "1h", ts.to_pydatetime(), *map(float, r * 1.001)) for ts, r in217 zip(df.index[:10], df.to_numpy()[:10])]218 rows += [("coinbase", "LINK", "1h", ts.to_pydatetime(), *map(float, r)) for ts, r in219 zip(df.index[::-1], df.to_numpy()[::-1])] # inserted newest first220 conn.cursor().executemany("INSERT INTO candles VALUES (%s,%s,%s,%s,%s,%s,%s,%s,%s)", rows)221222 out, q = load_candles("ETH", "1h", "2024-01-01T05:00", "2024-01-01T15:00", database_url=pg_url,223 source="coinbase", as_of=AS_OF)224 assert len(out) == 10 and out.index[0] == pd.Timestamp("2024-01-01T05:00", tz="UTC") and q.source == "coinbase"225 np.testing.assert_allclose(out.to_numpy(), df.iloc[5:15].to_numpy(), rtol=1e-12)226 assert out.dtypes.eq("float64").all() and out.index.is_monotonic_increasing227228 with pytest.raises(DataError, match="2 sources"): # two sources and none chosen: refuse to guess229 load_candles("ETH", "1h", database_url=pg_url, as_of=AS_OF)230 one, q1 = load_candles("LINK", "1h", database_url=pg_url, as_of=AS_OF) # a single source is unambiguous231 assert len(one) == 30 and q1.source == "coinbase" and one.index.is_monotonic_increasing232 with pytest.raises(DataError, match="no rows"):233 load_candles("ETH", "1h", "2030-01-01", database_url=pg_url, source="coinbase")234 # only closed bars leave the loader, even if the table holds a bar that is still open235 out, q = load_candles("ETH", "1h", database_url=pg_url, source="coinbase", as_of="2024-01-01T10:30")236 assert out.index[-1] == pd.Timestamp("2024-01-01T09:00", tz="UTC") and q.unclosed_dropped == 20237 # the session is read-only238 with psycopg.connect(pg_url, options="-c default_transaction_read_only=on") as conn:239 with pytest.raises(psycopg.errors.ReadOnlySqlTransaction):240 conn.execute("DELETE FROM candles")241