Unit tests cover the SQL gate, broadcast codec, Hive writer layout, and flight kinematics. CI runs pytest first; the simulator image build runs tests in a dedicated stage before the final layer.
68 lines
2.3 KiB
Python
68 lines
2.3 KiB
Python
"""Hive-partitioned Parquet writer layout."""
|
|
|
|
from __future__ import annotations
|
|
|
|
import time
|
|
from pathlib import Path
|
|
|
|
import duckdb
|
|
|
|
from virtual_drone.flight import Pose
|
|
from virtual_drone.sensors import imu_row
|
|
from virtual_drone.writer import PartitionWriter, make_writers
|
|
|
|
|
|
def test_partition_path_follows_hive_layout(tmp_path: Path) -> None:
|
|
ts_ns = int(time.time() * 1e9)
|
|
writer = PartitionWriter(tmp_path, "telemetry", "flt-01", "dr-01", "imu")
|
|
writer.append(imu_row(
|
|
Pose(ts_ns=ts_ns, x=1.0, y=2.0, z=3.0, roll=0.0, pitch=0.0, yaw=0.0, vx=0.0, vy=0.0, vz=0.0),
|
|
__import__("random").Random(0),
|
|
ts_ns,
|
|
))
|
|
writer.seal()
|
|
|
|
files = list(tmp_path.rglob("*.parquet"))
|
|
assert len(files) >= 1
|
|
path_str = str(files[0])
|
|
assert "dataset=telemetry" in path_str
|
|
assert "flight=flt-01" in path_str
|
|
assert "drone=dr-01" in path_str
|
|
assert "sensor=imu" in path_str
|
|
assert "year=" in path_str and "hour=" in path_str
|
|
|
|
|
|
def test_seal_compacts_current_blocks(tmp_path: Path) -> None:
|
|
writer = PartitionWriter(tmp_path, "state", "flt-02", "dr-02", None)
|
|
base_ts = 1_700_000_000_000_000_000
|
|
for i in range(3):
|
|
writer.append({"ts_ns": base_ts + i * 5_000_000_000, "pos_x": float(i)})
|
|
writer.seal()
|
|
|
|
sealed = list(tmp_path.rglob("data.parquet"))
|
|
assert len(sealed) == 1
|
|
assert not list(tmp_path.rglob("current/min_*.parquet"))
|
|
rows = duckdb.sql(f"SELECT count(*) FROM read_parquet('{sealed[0]}')").fetchone()[0]
|
|
assert rows == 3
|
|
|
|
|
|
def test_make_writers_registers_expected_datasets(tmp_path: Path) -> None:
|
|
writers = make_writers(tmp_path, "flt-03", "dr-03", ["imu", "battery"])
|
|
assert set(writers) == {"imu", "battery", "detections", "state"}
|
|
|
|
|
|
def test_duckdb_reads_written_partition(tmp_path: Path) -> None:
|
|
writers = make_writers(tmp_path, "flt-04", "dr-04", ["imu"])
|
|
ts_ns = int(time.time() * 1e9)
|
|
pose = Pose(ts_ns=ts_ns, x=0.0, y=0.0, z=0.0, roll=0.0, pitch=0.0, yaw=0.0, vx=0.0, vy=0.0, vz=0.0)
|
|
writers["imu"].append(imu_row(pose, __import__("random").Random(1), ts_ns))
|
|
for w in writers.values():
|
|
w.seal()
|
|
|
|
count = duckdb.sql(f"""
|
|
SELECT count(*) FROM read_parquet(
|
|
'{tmp_path}/dataset=telemetry/**/*.parquet',
|
|
hive_partitioning=true, union_by_name=true)
|
|
""").fetchone()[0]
|
|
assert count >= 1
|