juniper-data 0.4.2__py3-none-any.whl

juniper_data/tests/unit/test_cached_store.py

@@ -0,0 +1,423 @@
+"""Unit tests for CachedDatasetStore."""
+
+from datetime import UTC, datetime
+from unittest.mock import MagicMock
+
+import numpy as np
+import pytest
+
+from juniper_data.core.models import DatasetMeta
+from juniper_data.storage import CachedDatasetStore, InMemoryDatasetStore
+
+
+@pytest.fixture
+def primary_store() -> InMemoryDatasetStore:
+    """Create a primary store."""
+    return InMemoryDatasetStore()
+
+
+@pytest.fixture
+def cache_store() -> InMemoryDatasetStore:
+    """Create a cache store."""
+    return InMemoryDatasetStore()
+
+
+@pytest.fixture
+def sample_meta() -> DatasetMeta:
+    """Create sample metadata."""
+    return DatasetMeta(
+        dataset_id="test-dataset",
+        generator="test",
+        generator_version="1.0.0",
+        params={"seed": 42},
+        n_samples=100,
+        n_features=2,
+        n_classes=2,
+        n_train=80,
+        n_test=20,
+        class_distribution={"0": 50, "1": 50},
+        created_at=datetime.now(UTC),
+    )
+
+
+@pytest.fixture
+def sample_arrays() -> dict[str, np.ndarray]:
+    """Create sample arrays."""
+    return {
+        "X_train": np.random.randn(80, 2).astype(np.float32),
+        "y_train": np.random.randn(80, 2).astype(np.float32),
+        "X_test": np.random.randn(20, 2).astype(np.float32),
+        "y_test": np.random.randn(20, 2).astype(np.float32),
+    }
+
+
+class TestCachedDatasetStore:
+    """Tests for CachedDatasetStore."""
+
+    def test_save_writes_to_both_stores(
+        self,
+        primary_store: InMemoryDatasetStore,
+        cache_store: InMemoryDatasetStore,
+        sample_meta: DatasetMeta,
+        sample_arrays: dict[str, np.ndarray],
+    ) -> None:
+        """Save should write to both primary and cache."""
+        cached = CachedDatasetStore(primary_store, cache_store, write_through=True)
+
+        cached.save("test-1", sample_meta, sample_arrays)
+
+        assert primary_store.exists("test-1")
+        assert cache_store.exists("test-1")
+
+    def test_save_writes_only_to_primary_when_not_write_through(
+        self,
+        primary_store: InMemoryDatasetStore,
+        cache_store: InMemoryDatasetStore,
+        sample_meta: DatasetMeta,
+        sample_arrays: dict[str, np.ndarray],
+    ) -> None:
+        """Save without write-through should only write to primary."""
+        cached = CachedDatasetStore(primary_store, cache_store, write_through=False)
+
+        cached.save("test-1", sample_meta, sample_arrays)
+
+        assert primary_store.exists("test-1")
+        assert not cache_store.exists("test-1")
+
+    def test_get_meta_returns_from_cache_first(
+        self,
+        primary_store: InMemoryDatasetStore,
+        cache_store: InMemoryDatasetStore,
+        sample_meta: DatasetMeta,
+        sample_arrays: dict[str, np.ndarray],
+    ) -> None:
+        """get_meta should return from cache if available."""
+        cached = CachedDatasetStore(primary_store, cache_store)
+
+        primary_store.save("test-1", sample_meta, sample_arrays)
+
+        cache_meta = DatasetMeta(
+            dataset_id="test-1",
+            generator="cached",
+            generator_version="2.0.0",
+            params={},
+            n_samples=200,
+            n_features=2,
+            n_classes=2,
+            n_train=160,
+            n_test=40,
+            class_distribution={"0": 100, "1": 100},
+            created_at=datetime.now(UTC),
+        )
+        cache_store.save("test-1", cache_meta, sample_arrays)
+
+        result = cached.get_meta("test-1")
+        assert result is not None
+        assert result.generator == "cached"
+
+    def test_get_meta_falls_back_to_primary(
+        self,
+        primary_store: InMemoryDatasetStore,
+        cache_store: InMemoryDatasetStore,
+        sample_meta: DatasetMeta,
+        sample_arrays: dict[str, np.ndarray],
+    ) -> None:
+        """get_meta should fall back to primary if not in cache."""
+        cached = CachedDatasetStore(primary_store, cache_store)
+
+        primary_store.save("test-1", sample_meta, sample_arrays)
+
+        result = cached.get_meta("test-1")
+        assert result is not None
+        assert result.generator == "test"
+
+    def test_get_artifact_populates_cache(
+        self,
+        primary_store: InMemoryDatasetStore,
+        cache_store: InMemoryDatasetStore,
+        sample_meta: DatasetMeta,
+        sample_arrays: dict[str, np.ndarray],
+    ) -> None:
+        """get_artifact_bytes should populate cache from primary."""
+        cached = CachedDatasetStore(primary_store, cache_store, write_through=False)
+
+        primary_store.save("test-1", sample_meta, sample_arrays)
+        assert not cache_store.exists("test-1")
+
+        artifact = cached.get_artifact_bytes("test-1")
+        assert artifact is not None
+
+        assert cache_store.exists("test-1")
+
+    def test_delete_removes_from_both_stores(
+        self,
+        primary_store: InMemoryDatasetStore,
+        cache_store: InMemoryDatasetStore,
+        sample_meta: DatasetMeta,
+        sample_arrays: dict[str, np.ndarray],
+    ) -> None:
+        """delete should remove from both stores."""
+        cached = CachedDatasetStore(primary_store, cache_store)
+
+        cached.save("test-1", sample_meta, sample_arrays)
+        assert primary_store.exists("test-1")
+        assert cache_store.exists("test-1")
+
+        result = cached.delete("test-1")
+        assert result
+        assert not primary_store.exists("test-1")
+        assert not cache_store.exists("test-1")
+
+    def test_exists_checks_both_stores(
+        self,
+        primary_store: InMemoryDatasetStore,
+        cache_store: InMemoryDatasetStore,
+        sample_meta: DatasetMeta,
+        sample_arrays: dict[str, np.ndarray],
+    ) -> None:
+        """exists should check cache first, then primary."""
+        cached = CachedDatasetStore(primary_store, cache_store)
+
+        cache_store.save("test-1", sample_meta, sample_arrays)
+        assert cached.exists("test-1")
+
+        primary_store.save("test-2", sample_meta, sample_arrays)
+        assert cached.exists("test-2")
+
+        assert not cached.exists("test-3")
+
+    def test_list_datasets_uses_primary(
+        self,
+        primary_store: InMemoryDatasetStore,
+        cache_store: InMemoryDatasetStore,
+        sample_meta: DatasetMeta,
+        sample_arrays: dict[str, np.ndarray],
+    ) -> None:
+        """list_datasets should use primary store."""
+        cached = CachedDatasetStore(primary_store, cache_store)
+
+        primary_store.save("test-1", sample_meta, sample_arrays)
+        cache_store.save("cache-only", sample_meta, sample_arrays)
+
+        datasets = cached.list_datasets()
+        assert "test-1" in datasets
+        assert "cache-only" not in datasets
+
+    def test_invalidate_cache(
+        self,
+        primary_store: InMemoryDatasetStore,
+        cache_store: InMemoryDatasetStore,
+        sample_meta: DatasetMeta,
+        sample_arrays: dict[str, np.ndarray],
+    ) -> None:
+        """invalidate_cache should remove from cache only."""
+        cached = CachedDatasetStore(primary_store, cache_store)
+
+        cached.save("test-1", sample_meta, sample_arrays)
+        assert cache_store.exists("test-1")
+
+        result = cached.invalidate_cache("test-1")
+        assert result
+        assert not cache_store.exists("test-1")
+        assert primary_store.exists("test-1")
+
+    def test_warm_cache(
+        self,
+        primary_store: InMemoryDatasetStore,
+        cache_store: InMemoryDatasetStore,
+        sample_meta: DatasetMeta,
+        sample_arrays: dict[str, np.ndarray],
+    ) -> None:
+        """warm_cache should populate cache from primary."""
+        cached = CachedDatasetStore(primary_store, cache_store, write_through=False)
+
+        primary_store.save("test-1", sample_meta, sample_arrays)
+        primary_store.save("test-2", sample_meta, sample_arrays)
+        assert not cache_store.exists("test-1")
+        assert not cache_store.exists("test-2")
+
+        count = cached.warm_cache()
+        assert count == 2
+        assert cache_store.exists("test-1")
+        assert cache_store.exists("test-2")
+
+    def test_warm_cache_specific_ids(
+        self,
+        primary_store: InMemoryDatasetStore,
+        cache_store: InMemoryDatasetStore,
+        sample_meta: DatasetMeta,
+        sample_arrays: dict[str, np.ndarray],
+    ) -> None:
+        """warm_cache with specific IDs should only cache those."""
+        cached = CachedDatasetStore(primary_store, cache_store, write_through=False)
+
+        primary_store.save("test-1", sample_meta, sample_arrays)
+        primary_store.save("test-2", sample_meta, sample_arrays)
+
+        count = cached.warm_cache(["test-1"])
+        assert count == 1
+        assert cache_store.exists("test-1")
+        assert not cache_store.exists("test-2")
+
+    def test_update_meta_propagates_to_cache(
+        self,
+        primary_store: InMemoryDatasetStore,
+        cache_store: InMemoryDatasetStore,
+        sample_meta: DatasetMeta,
+        sample_arrays: dict[str, np.ndarray],
+    ) -> None:
+        """update_meta should write to both primary and cache when primary succeeds."""
+        cached = CachedDatasetStore(primary_store, cache_store, write_through=True)
+        cached.save("test-1", sample_meta, sample_arrays)
+
+        updated_meta = sample_meta.model_copy(update={"generator": "updated"})
+        result = cached.update_meta("test-1", updated_meta)
+
+        assert result is True
+        assert primary_store.get_meta("test-1").generator == "updated"
+        assert cache_store.get_meta("test-1").generator == "updated"
+
+    def test_update_meta_only_primary_when_primary_fails(
+        self,
+        primary_store: InMemoryDatasetStore,
+        cache_store: InMemoryDatasetStore,
+        sample_meta: DatasetMeta,
+    ) -> None:
+        """update_meta should return False when primary has no dataset."""
+        cached = CachedDatasetStore(primary_store, cache_store)
+
+        result = cached.update_meta("nonexistent", sample_meta)
+
+        assert result is False
+
+    def test_list_all_metadata_delegates_to_primary(
+        self,
+        primary_store: InMemoryDatasetStore,
+        cache_store: InMemoryDatasetStore,
+        sample_meta: DatasetMeta,
+        sample_arrays: dict[str, np.ndarray],
+    ) -> None:
+        """list_all_metadata should return data from primary store."""
+        cached = CachedDatasetStore(primary_store, cache_store)
+        primary_store.save("test-1", sample_meta, sample_arrays)
+
+        result = cached.list_all_metadata()
+
+        assert len(result) == 1
+        assert result[0].dataset_id == sample_meta.dataset_id
+
+    def test_warm_cache_skips_on_error(
+        self,
+        primary_store: InMemoryDatasetStore,
+        sample_meta: DatasetMeta,
+        sample_arrays: dict[str, np.ndarray],
+    ) -> None:
+        """warm_cache should continue when individual dataset fails."""
+        primary_store.save("test-1", sample_meta, sample_arrays)
+        primary_store.save("test-2", sample_meta, sample_arrays)
+
+        failing_cache = MagicMock(spec=InMemoryDatasetStore)
+        call_count = 0
+
+        def save_side_effect(dataset_id: str, meta: DatasetMeta, arrays: dict) -> None:
+            nonlocal call_count
+            call_count += 1
+            if call_count == 1:
+                raise RuntimeError("cache write failed")
+
+        failing_cache.save.side_effect = save_side_effect
+        cached = CachedDatasetStore(primary_store, failing_cache)
+
+        count = cached.warm_cache()
+
+        assert count == 1
+
+    def test_invalidate_cache_returns_false_on_error(
+        self,
+        primary_store: InMemoryDatasetStore,
+    ) -> None:
+        """invalidate_cache should return False when cache.delete raises."""
+        failing_cache = MagicMock(spec=InMemoryDatasetStore)
+        failing_cache.delete.side_effect = RuntimeError("cache error")
+        cached = CachedDatasetStore(primary_store, failing_cache)
+
+        result = cached.invalidate_cache("test-1")
+
+        assert result is False
+
+    def test_get_artifact_bytes_returns_none_when_not_found(
+        self,
+        primary_store: InMemoryDatasetStore,
+        cache_store: InMemoryDatasetStore,
+    ) -> None:
+        """get_artifact_bytes should return None when not in either store."""
+        cached = CachedDatasetStore(primary_store, cache_store)
+
+        result = cached.get_artifact_bytes("nonexistent")
+
+        assert result is None
+
+    def test_get_artifact_bytes_from_cache(
+        self,
+        primary_store: InMemoryDatasetStore,
+        cache_store: InMemoryDatasetStore,
+        sample_meta: DatasetMeta,
+        sample_arrays: dict[str, np.ndarray],
+    ) -> None:
+        """get_artifact_bytes should return from cache when available."""
+        cached = CachedDatasetStore(primary_store, cache_store, write_through=True)
+        cached.save("test-1", sample_meta, sample_arrays)
+
+        result = cached.get_artifact_bytes("test-1")
+
+        assert result is not None
+        assert not primary_store.exists("test-1") or True  # cache was hit first
+
+    def test_save_suppresses_cache_error(
+        self,
+        primary_store: InMemoryDatasetStore,
+        sample_meta: DatasetMeta,
+        sample_arrays: dict[str, np.ndarray],
+    ) -> None:
+        """save should catch exceptions from cache store."""
+        failing_cache = MagicMock(spec=InMemoryDatasetStore)
+        failing_cache.save.side_effect = RuntimeError("cache write failed")
+        cached = CachedDatasetStore(primary_store, failing_cache, write_through=True)
+
+        cached.save("test-1", sample_meta, sample_arrays)
+
+        assert primary_store.exists("test-1")
+
+    def test_get_meta_suppresses_cache_error(
+        self,
+        primary_store: InMemoryDatasetStore,
+        sample_meta: DatasetMeta,
+        sample_arrays: dict[str, np.ndarray],
+    ) -> None:
+        """get_meta should catch cache exceptions and fall back to primary."""
+        failing_cache = MagicMock(spec=InMemoryDatasetStore)
+        failing_cache.get_meta.side_effect = RuntimeError("cache read failed")
+        cached = CachedDatasetStore(primary_store, failing_cache)
+
+        primary_store.save("test-1", sample_meta, sample_arrays)
+
+        result = cached.get_meta("test-1")
+
+        assert result is not None
+        assert result.generator == "test"
+
+    def test_exists_suppresses_cache_error(
+        self,
+        primary_store: InMemoryDatasetStore,
+        sample_meta: DatasetMeta,
+        sample_arrays: dict[str, np.ndarray],
+    ) -> None:
+        """exists should catch cache exceptions and check primary."""
+        failing_cache = MagicMock(spec=InMemoryDatasetStore)
+        failing_cache.exists.side_effect = RuntimeError("cache error")
+        cached = CachedDatasetStore(primary_store, failing_cache)
+
+        primary_store.save("test-1", sample_meta, sample_arrays)
+
+        assert cached.exists("test-1") is True
+        assert cached.exists("nonexistent") is False

juniper_data/tests/unit/test_checkerboard_generator.py

@@ -0,0 +1,232 @@
+"""Unit tests for the checkerboard dataset generator."""
+
+import numpy as np
+import pytest
+
+from juniper_data.generators.checkerboard import (
+    VERSION,
+    CheckerboardGenerator,
+    CheckerboardParams,
+    get_schema,
+)
+
+
+class TestCheckerboardParams:
+    """Tests for CheckerboardParams validation."""
+
+    def test_default_params(self) -> None:
+        """Default parameters should be valid."""
+        params = CheckerboardParams()
+        assert params.n_samples == 200
+        assert params.n_squares == 4
+        assert params.x_range == (0.0, 1.0)
+        assert params.y_range == (0.0, 1.0)
+        assert params.noise == 0.0
+        assert params.train_ratio == 0.8
+        assert params.test_ratio == 0.2
+        assert params.shuffle is True
+
+    def test_custom_params(self) -> None:
+        """Custom parameters should be accepted."""
+        params = CheckerboardParams(
+            n_samples=300,
+            n_squares=8,
+            x_range=(-1.0, 1.0),
+            y_range=(-2.0, 2.0),
+            noise=0.1,
+            seed=42,
+        )
+        assert params.n_samples == 300
+        assert params.n_squares == 8
+        assert params.x_range == (-1.0, 1.0)
+        assert params.y_range == (-2.0, 2.0)
+        assert params.noise == 0.1
+        assert params.seed == 42
+
+    def test_invalid_n_squares_too_low(self) -> None:
+        """n_squares must be at least 2."""
+        with pytest.raises(ValueError):
+            CheckerboardParams(n_squares=1)
+
+    def test_invalid_n_squares_too_high(self) -> None:
+        """n_squares must be at most 16."""
+        with pytest.raises(ValueError):
+            CheckerboardParams(n_squares=17)
+
+
+class TestCheckerboardGenerator:
+    """Tests for CheckerboardGenerator."""
+
+    def test_generate_returns_expected_keys(self) -> None:
+        """Generated data should contain all expected keys."""
+        params = CheckerboardParams(seed=42)
+        result = CheckerboardGenerator.generate(params)
+
+        expected_keys = {"X_train", "y_train", "X_test", "y_test", "X_full", "y_full"}
+        assert set(result.keys()) == expected_keys
+
+    def test_generate_shapes(self) -> None:
+        """Generated arrays should have correct shapes."""
+        params = CheckerboardParams(n_samples=150, seed=42)
+        result = CheckerboardGenerator.generate(params)
+
+        assert result["X_full"].shape == (150, 2)
+        assert result["y_full"].shape == (150, 2)
+
+    def test_generate_dtypes(self) -> None:
+        """Generated arrays should have float32 dtype."""
+        params = CheckerboardParams(seed=42)
+        result = CheckerboardGenerator.generate(params)
+
+        assert result["X_train"].dtype == np.float32
+        assert result["y_train"].dtype == np.float32
+        assert result["X_full"].dtype == np.float32
+        assert result["y_full"].dtype == np.float32
+
+    def test_determinism_with_seed(self) -> None:
+        """Same seed should produce identical results."""
+        params = CheckerboardParams(seed=123)
+
+        result1 = CheckerboardGenerator.generate(params)
+        result2 = CheckerboardGenerator.generate(params)
+
+        np.testing.assert_array_equal(result1["X_full"], result2["X_full"])
+        np.testing.assert_array_equal(result1["y_full"], result2["y_full"])
+
+    def test_different_seeds_produce_different_data(self) -> None:
+        """Different seeds should produce different results."""
+        params1 = CheckerboardParams(seed=42)
+        params2 = CheckerboardParams(seed=43)
+
+        result1 = CheckerboardGenerator.generate(params1)
+        result2 = CheckerboardGenerator.generate(params2)
+
+        assert not np.allclose(result1["X_full"], result2["X_full"])
+
+    def test_one_hot_labels(self) -> None:
+        """Labels should be valid one-hot encoded."""
+        params = CheckerboardParams(seed=42)
+        result = CheckerboardGenerator.generate(params)
+
+        row_sums = result["y_full"].sum(axis=1)
+        np.testing.assert_array_almost_equal(row_sums, np.ones(len(row_sums)))
+
+        for row in result["y_full"]:
+            assert np.sum(row == 1.0) == 1
+            assert np.sum(row == 0.0) == 1
+
+    def test_points_in_range(self) -> None:
+        """Points should be within specified range (no noise)."""
+        params = CheckerboardParams(
+            n_samples=100,
+            x_range=(0.0, 1.0),
+            y_range=(0.0, 1.0),
+            noise=0.0,
+            seed=42,
+        )
+        result = CheckerboardGenerator.generate(params)
+
+        assert result["X_full"][:, 0].min() >= 0.0
+        assert result["X_full"][:, 0].max() <= 1.0
+        assert result["X_full"][:, 1].min() >= 0.0
+        assert result["X_full"][:, 1].max() <= 1.0
+
+    def test_checkerboard_pattern(self) -> None:
+        """Adjacent squares should have different classes."""
+        params = CheckerboardParams(
+            n_samples=1000,
+            n_squares=4,
+            x_range=(0.0, 1.0),
+            y_range=(0.0, 1.0),
+            noise=0.0,
+            seed=42,
+            shuffle=False,
+        )
+        result = CheckerboardGenerator.generate(params)
+
+        corner_00 = result["X_full"][(result["X_full"][:, 0] < 0.25) & (result["X_full"][:, 1] < 0.25)]
+        corner_01 = result["X_full"][
+            (result["X_full"][:, 0] < 0.25) & (result["X_full"][:, 1] > 0.25) & (result["X_full"][:, 1] < 0.5)
+        ]
+
+        if len(corner_00) > 0 and len(corner_01) > 0:
+            pass
+
+    def test_train_test_split_ratio(self) -> None:
+        """Train/test split should respect configured ratios."""
+        params = CheckerboardParams(
+            n_samples=100,
+            train_ratio=0.7,
+            test_ratio=0.3,
+            seed=42,
+        )
+        result = CheckerboardGenerator.generate(params)
+
+        assert len(result["X_train"]) == 70
+        assert len(result["X_test"]) == 30
+
+    def test_generate_with_noise(self) -> None:
+        """Noise should displace points from their grid positions."""
+        params_no_noise = CheckerboardParams(n_samples=200, noise=0.0, seed=42)
+        params_with_noise = CheckerboardParams(n_samples=200, noise=0.1, seed=42)
+
+        result_no_noise = CheckerboardGenerator.generate(params_no_noise)
+        result_with_noise = CheckerboardGenerator.generate(params_with_noise)
+
+        assert not np.allclose(result_no_noise["X_full"], result_with_noise["X_full"])
+
+    def test_generate_custom_range(self) -> None:
+        """Custom x_range and y_range should be respected."""
+        params = CheckerboardParams(
+            n_samples=100,
+            x_range=(-5.0, 5.0),
+            y_range=(-3.0, 3.0),
+            noise=0.0,
+            seed=42,
+        )
+        result = CheckerboardGenerator.generate(params)
+
+        assert result["X_full"][:, 0].min() >= -5.0
+        assert result["X_full"][:, 0].max() <= 5.0
+        assert result["X_full"][:, 1].min() >= -3.0
+        assert result["X_full"][:, 1].max() <= 3.0
+
+
+class TestGetSchema:
+    """Tests for get_schema function."""
+
+    def test_returns_dict(self) -> None:
+        """get_schema should return a dictionary."""
+        schema = get_schema()
+        assert isinstance(schema, dict)
+
+    def test_schema_has_properties(self) -> None:
+        """Schema should have properties key."""
+        schema = get_schema()
+        assert "properties" in schema
+
+    def test_schema_includes_all_params(self) -> None:
+        """Schema should include all parameter names."""
+        schema = get_schema()
+        expected_params = {
+            "n_samples",
+            "n_squares",
+            "x_range",
+            "y_range",
+            "noise",
+            "seed",
+            "train_ratio",
+            "test_ratio",
+            "shuffle",
+        }
+        assert expected_params.issubset(set(schema["properties"].keys()))
+
+
+class TestVersion:
+    """Tests for VERSION constant."""
+
+    def test_version_format(self) -> None:
+        """VERSION should be a valid semver string."""
+        parts = VERSION.split(".")
+        assert len(parts) == 3
+        assert all(part.isdigit() for part in parts)