orderflow_backtest/tests/test_repository_iter.py

54 lines
1.7 KiB
Python

import sys
from pathlib import Path
import sqlite3
sys.path.append(str(Path(__file__).resolve().parents[1]))
from repositories.sqlite_repository import SQLiteOrderflowRepository
def test_iterate_book_rows_batches(tmp_path):
db_path = tmp_path / "iter.db"
with sqlite3.connect(str(db_path)) as conn:
c = conn.cursor()
c.execute(
"""
CREATE TABLE book (
id INTEGER PRIMARY KEY,
bids TEXT NOT NULL,
asks TEXT NOT NULL,
timestamp INTEGER NOT NULL
)
"""
)
c.execute(
"""
CREATE TABLE trades (
id INTEGER PRIMARY KEY,
trade_id REAL NOT NULL,
price REAL NOT NULL,
size REAL NOT NULL,
side TEXT NOT NULL,
timestamp INTEGER NOT NULL
)
"""
)
# Insert 12 rows to ensure multiple fetchmany batches (repo uses 5000, but iteration still correct)
bids = str([(100.0, 1.0, 0, 1)])
asks = str([(101.0, 1.0, 0, 1)])
for i in range(12):
c.execute("INSERT INTO book (id, bids, asks, timestamp) VALUES (?, ?, ?, ?)", (i + 1, bids, asks, 1000 + i))
conn.commit()
repo = SQLiteOrderflowRepository(db_path)
with repo.connect() as conn:
rows = list(repo.iterate_book_rows(conn))
assert len(rows) == 12
# Ensure ordering by timestamp ascending
timestamps = [r[3] for r in rows]
assert timestamps == sorted(timestamps)
# count_rows allowlist should work
assert repo.count_rows(conn, "book") == 12