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.
85 lines
3.3 KiB
Python
85 lines
3.3 KiB
Python
"""Unit tests for snapshot export."""
|
|
from __future__ import annotations
|
|
|
|
import json
|
|
import tempfile
|
|
from pathlib import Path
|
|
from unittest.mock import AsyncMock, MagicMock
|
|
|
|
import pytest
|
|
|
|
from libs.export.snapshot_export import _temporal_split, export_dataset_snapshot
|
|
|
|
|
|
@pytest.mark.unit
|
|
class TestTemporalSplit:
|
|
def test_split_proportions(self) -> None:
|
|
"""Temporal split produces correct proportions for 100 rows."""
|
|
rows = [{"event_date": "2026-01-01"} for _ in range(100)]
|
|
splits = _temporal_split(rows, "temporal_70_15_15")
|
|
assert len(splits["train"]) == 70
|
|
assert len(splits["valid"]) == 15
|
|
assert len(splits["test"]) == 15
|
|
|
|
def test_split_preserves_temporal_order(self) -> None:
|
|
"""Train set contains earliest dates, test contains latest."""
|
|
rows = [{"event_date": f"2026-{m:02d}-01"} for m in range(1, 13)]
|
|
splits = _temporal_split(rows, "temporal_70_15_15")
|
|
if splits["train"] and splits["test"]:
|
|
assert splits["train"][-1]["event_date"] <= splits["test"][0]["event_date"]
|
|
|
|
def test_empty_rows_returns_empty_splits(self) -> None:
|
|
splits = _temporal_split([], "temporal_70_15_15")
|
|
assert splits == {"train": [], "valid": [], "test": []}
|
|
|
|
|
|
@pytest.mark.unit
|
|
class TestExportDatasetSnapshot:
|
|
@pytest.mark.asyncio
|
|
async def test_manifest_is_written(self) -> None:
|
|
"""export_dataset_snapshot writes a manifest.json with expected fields."""
|
|
# Mock DB session returning empty results (no feature+label pairs)
|
|
mock_session = AsyncMock()
|
|
mock_result = MagicMock()
|
|
mock_result.all.return_value = [] # no rows
|
|
mock_session.execute = AsyncMock(return_value=mock_result)
|
|
|
|
with tempfile.TemporaryDirectory() as tmpdir:
|
|
manifest = await export_dataset_snapshot(
|
|
session=mock_session,
|
|
snapshot_id="test-snapshot-001",
|
|
split_policy="temporal_70_15_15",
|
|
output_dir=tmpdir,
|
|
)
|
|
|
|
assert manifest["snapshot_id"] == "test-snapshot-001"
|
|
assert "created_at_utc" in manifest
|
|
assert "row_counts" in manifest
|
|
assert manifest["split_policy"] == "temporal_70_15_15"
|
|
assert manifest["total_rows"] == 0
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_parquet_files_created(self) -> None:
|
|
"""Parquet files are created for each split partition."""
|
|
mock_session = AsyncMock()
|
|
mock_result = MagicMock()
|
|
mock_result.all.return_value = []
|
|
mock_session.execute = AsyncMock(return_value=mock_result)
|
|
|
|
with tempfile.TemporaryDirectory() as tmpdir:
|
|
await export_dataset_snapshot(
|
|
session=mock_session,
|
|
snapshot_id="test-parquet-002",
|
|
split_policy="temporal_70_15_15",
|
|
output_dir=tmpdir,
|
|
)
|
|
|
|
snap_dir = Path(tmpdir) / "test-parquet-002"
|
|
assert (snap_dir / "train.parquet").exists()
|
|
assert (snap_dir / "valid.parquet").exists()
|
|
assert (snap_dir / "test.parquet").exists()
|
|
assert (snap_dir / "manifest.json").exists()
|
|
|
|
manifest_data = json.loads((snap_dir / "manifest.json").read_text())
|
|
assert manifest_data["snapshot_id"] == "test-parquet-002"
|