juniper-data 0.4.2__py3-none-any.whl

juniper_data/tests/integration/test_security_integration.py

@@ -0,0 +1,189 @@
+"""Integration tests for API security middleware."""
+
+import pytest
+from fastapi.testclient import TestClient
+
+from juniper_data.api.app import create_app
+from juniper_data.api.routes import datasets
+from juniper_data.api.settings import Settings
+from juniper_data.storage import InMemoryDatasetStore
+
+
+@pytest.fixture
+def auth_enabled_client() -> TestClient:
+    """Create a test client with API key authentication enabled."""
+    settings = Settings(
+        storage_path="./test_data",
+        api_keys=["valid-key-1", "valid-key-2"],
+        rate_limit_enabled=False,
+    )
+    app = create_app(settings)
+    datasets.set_store(InMemoryDatasetStore())
+    return TestClient(app)
+
+
+@pytest.fixture
+def rate_limited_client() -> TestClient:
+    """Create a test client with rate limiting enabled."""
+    settings = Settings(
+        storage_path="./test_data",
+        api_keys=None,
+        rate_limit_enabled=True,
+        rate_limit_requests_per_minute=5,
+    )
+    app = create_app(settings)
+    datasets.set_store(InMemoryDatasetStore())
+    return TestClient(app)
+
+
+@pytest.fixture
+def fully_secured_client() -> TestClient:
+    """Create a test client with both auth and rate limiting enabled."""
+    settings = Settings(
+        storage_path="./test_data",
+        api_keys=["secure-key"],
+        rate_limit_enabled=True,
+        rate_limit_requests_per_minute=10,
+    )
+    app = create_app(settings)
+    datasets.set_store(InMemoryDatasetStore())
+    return TestClient(app)
+
+
+class TestAPIKeyAuthentication:
+    """Integration tests for API key authentication."""
+
+    def test_health_endpoint_exempt(self, auth_enabled_client: TestClient) -> None:
+        """Health endpoints should be accessible without API key."""
+        response = auth_enabled_client.get("/v1/health")
+        assert response.status_code == 200
+
+        response = auth_enabled_client.get("/v1/health/live")
+        assert response.status_code == 200
+
+        response = auth_enabled_client.get("/v1/health/ready")
+        assert response.status_code == 200
+
+    def test_protected_endpoint_requires_key(self, auth_enabled_client: TestClient) -> None:
+        """Protected endpoints should require API key."""
+        response = auth_enabled_client.get("/v1/generators")
+        assert response.status_code == 401
+        assert "Missing API key" in response.json()["detail"]
+
+    def test_invalid_key_rejected(self, auth_enabled_client: TestClient) -> None:
+        """Invalid API keys should be rejected."""
+        response = auth_enabled_client.get(
+            "/v1/generators",
+            headers={"X-API-Key": "invalid-key"},
+        )
+        assert response.status_code == 401
+        assert "Invalid API key" in response.json()["detail"]
+
+    def test_valid_key_accepted(self, auth_enabled_client: TestClient) -> None:
+        """Valid API keys should be accepted."""
+        response = auth_enabled_client.get(
+            "/v1/generators",
+            headers={"X-API-Key": "valid-key-1"},
+        )
+        assert response.status_code == 200
+
+        response = auth_enabled_client.get(
+            "/v1/generators",
+            headers={"X-API-Key": "valid-key-2"},
+        )
+        assert response.status_code == 200
+
+    def test_create_dataset_with_auth(self, auth_enabled_client: TestClient) -> None:
+        """Dataset creation should work with valid API key."""
+        response = auth_enabled_client.post(
+            "/v1/datasets",
+            json={"generator": "spiral", "params": {"seed": 42}},
+            headers={"X-API-Key": "valid-key-1"},
+        )
+        assert response.status_code == 201
+        assert "dataset_id" in response.json()
+
+
+class TestRateLimiting:
+    """Integration tests for rate limiting."""
+
+    def test_health_endpoint_exempt(self, rate_limited_client: TestClient) -> None:
+        """Health endpoints should be exempt from rate limiting."""
+        for _ in range(20):
+            response = rate_limited_client.get("/v1/health")
+            assert response.status_code == 200
+
+    def test_allows_requests_within_limit(self, rate_limited_client: TestClient) -> None:
+        """Requests within limit should be allowed."""
+        for i in range(5):
+            response = rate_limited_client.get("/v1/generators")
+            assert response.status_code == 200
+            assert "X-RateLimit-Remaining" in response.headers
+            assert int(response.headers["X-RateLimit-Remaining"]) == 4 - i
+
+    def test_blocks_requests_over_limit(self, rate_limited_client: TestClient) -> None:
+        """Requests over limit should be blocked with 429."""
+        for _ in range(5):
+            rate_limited_client.get("/v1/generators")
+
+        response = rate_limited_client.get("/v1/generators")
+        assert response.status_code == 429
+        assert "Rate limit exceeded" in response.json()["detail"]
+        assert "Retry-After" in response.headers
+
+    def test_rate_limit_headers_present(self, rate_limited_client: TestClient) -> None:
+        """Rate limit headers should be present in responses."""
+        response = rate_limited_client.get("/v1/generators")
+        assert response.status_code == 200
+
+        assert "X-RateLimit-Limit" in response.headers
+        assert response.headers["X-RateLimit-Limit"] == "5"
+        assert "X-RateLimit-Remaining" in response.headers
+        assert "X-RateLimit-Reset" in response.headers
+
+
+class TestCombinedSecurity:
+    """Integration tests for combined auth and rate limiting."""
+
+    def test_auth_checked_before_rate_limit(self, fully_secured_client: TestClient) -> None:
+        """Authentication should be checked before rate limiting."""
+        response = fully_secured_client.get("/v1/generators")
+        assert response.status_code == 401
+
+        response = fully_secured_client.get(
+            "/v1/generators",
+            headers={"X-API-Key": "invalid"},
+        )
+        assert response.status_code == 401
+
+    def test_rate_limit_applied_after_auth(self, fully_secured_client: TestClient) -> None:
+        """Rate limiting should be applied after successful auth."""
+        for _ in range(10):
+            response = fully_secured_client.get(
+                "/v1/generators",
+                headers={"X-API-Key": "secure-key"},
+            )
+            assert response.status_code == 200
+
+        response = fully_secured_client.get(
+            "/v1/generators",
+            headers={"X-API-Key": "secure-key"},
+        )
+        assert response.status_code == 429
+
+    def test_full_workflow_with_security(self, fully_secured_client: TestClient) -> None:
+        """Full dataset workflow should work with security enabled."""
+        response = fully_secured_client.post(
+            "/v1/datasets",
+            json={"generator": "xor", "params": {"seed": 42}},
+            headers={"X-API-Key": "secure-key"},
+        )
+        assert response.status_code == 201
+        dataset_id = response.json()["dataset_id"]
+
+        response = fully_secured_client.get(
+            f"/v1/datasets/{dataset_id}",
+            headers={"X-API-Key": "secure-key"},
+        )
+        assert response.status_code == 200
+        assert response.json()["dataset_id"] == dataset_id

juniper_data/tests/integration/test_storage_workflow.py

@@ -0,0 +1,259 @@
+"""Integration tests for storage workflows.
+
+Tests cover end-to-end scenarios:
+- Generate → Store → Retrieve → Verify
+- Cross-store compatibility
+- Full dataset lifecycle
+"""
+
+import io
+import tempfile
+from datetime import datetime
+from pathlib import Path
+
+import numpy as np
+import pytest
+
+from juniper_data.core.artifacts import compute_checksum
+from juniper_data.core.models import DatasetMeta
+from juniper_data.generators.spiral import SpiralGenerator, SpiralParams
+from juniper_data.storage import InMemoryDatasetStore, LocalFSDatasetStore
+
+
+@pytest.fixture
+def temp_storage_dir():
+    """Create a temporary directory for storage tests."""
+    with tempfile.TemporaryDirectory() as tmpdir:
+        yield Path(tmpdir)
+
+
+@pytest.fixture
+def spiral_params() -> SpiralParams:
+    """Standard spiral parameters for testing."""
+    return SpiralParams(
+        n_spirals=2,
+        n_points_per_spiral=100,
+        noise=0.1,
+        seed=42,
+    )
+
+
+def create_dataset_meta(
+    dataset_id: str, params: SpiralParams, X: np.ndarray, y: np.ndarray, n_train: int, n_test: int
+) -> DatasetMeta:
+    """Helper to create DatasetMeta from generated data."""
+    n_classes = y.shape[1] if len(y.shape) > 1 else len(np.unique(y))
+    class_counts = np.sum(y, axis=0).astype(int) if len(y.shape) > 1 else np.bincount(y.astype(int))
+
+    return DatasetMeta(
+        dataset_id=dataset_id,
+        generator="spiral",
+        generator_version="1.0.0",
+        params=params.model_dump(),
+        n_samples=len(X),
+        n_features=X.shape[1],
+        n_classes=n_classes,
+        n_train=n_train,
+        n_test=n_test,
+        class_distribution={str(i): int(c) for i, c in enumerate(class_counts)},
+        created_at=datetime.now(),
+    )
+
+
+@pytest.mark.integration
+class TestGenerateStoreRetrieveWorkflow:
+    """Tests for the complete generate → store → retrieve workflow."""
+
+    def test_memory_store_full_workflow(self, spiral_params: SpiralParams):
+        """Generate spiral → store in memory → retrieve and verify."""
+        store = InMemoryDatasetStore()
+
+        data = SpiralGenerator.generate(spiral_params)
+        X, y = data["X_full"], data["y_full"]
+        X_train, y_train = data["X_train"], data["y_train"]
+        X_test, y_test = data["X_test"], data["y_test"]
+
+        arrays = {"X_train": X_train, "y_train": y_train, "X_test": X_test, "y_test": y_test}
+
+        meta = create_dataset_meta("spiral-mem-001", spiral_params, X, y, len(X_train), len(X_test))
+
+        store.save("spiral-mem-001", meta, arrays)
+
+        assert store.exists("spiral-mem-001")
+        retrieved_meta = store.get_meta("spiral-mem-001")
+        assert retrieved_meta is not None
+        assert retrieved_meta.dataset_id == "spiral-mem-001"
+        assert retrieved_meta.n_samples == 200
+
+        artifact_bytes = store.get_artifact_bytes("spiral-mem-001")
+        assert artifact_bytes is not None
+        loaded = np.load(io.BytesIO(artifact_bytes))
+        np.testing.assert_array_equal(loaded["X_train"], X_train)
+        np.testing.assert_array_equal(loaded["y_train"], y_train)
+
+    def test_fs_store_full_workflow(self, temp_storage_dir: Path, spiral_params: SpiralParams):
+        """Generate spiral → store to filesystem → retrieve and verify."""
+        store = LocalFSDatasetStore(temp_storage_dir)
+
+        data = SpiralGenerator.generate(spiral_params)
+        X, y = data["X_full"], data["y_full"]
+        X_train, y_train = data["X_train"], data["y_train"]
+        X_test, y_test = data["X_test"], data["y_test"]
+
+        arrays = {"X_train": X_train, "y_train": y_train, "X_test": X_test, "y_test": y_test}
+
+        meta = create_dataset_meta("spiral-fs-001", spiral_params, X, y, len(X_train), len(X_test))
+
+        store.save("spiral-fs-001", meta, arrays)
+
+        assert store.exists("spiral-fs-001")
+        assert (temp_storage_dir / "spiral-fs-001.meta.json").exists()
+        assert (temp_storage_dir / "spiral-fs-001.npz").exists()
+
+        retrieved_meta = store.get_meta("spiral-fs-001")
+        assert retrieved_meta is not None
+        assert retrieved_meta.dataset_id == "spiral-fs-001"
+        assert retrieved_meta.n_samples == 200
+
+        artifact_bytes = store.get_artifact_bytes("spiral-fs-001")
+        assert artifact_bytes is not None
+        loaded = np.load(io.BytesIO(artifact_bytes))
+        np.testing.assert_array_equal(loaded["X_train"], X_train)
+
+    def test_persistence_across_store_instances(self, temp_storage_dir: Path, spiral_params: SpiralParams):
+        """Data persists when creating new store instance."""
+        store1 = LocalFSDatasetStore(temp_storage_dir)
+
+        data = SpiralGenerator.generate(spiral_params)
+        X, y = data["X_full"], data["y_full"]
+        X_train, y_train = data["X_train"], data["y_train"]
+        X_test, y_test = data["X_test"], data["y_test"]
+
+        arrays = {"X_train": X_train, "y_train": y_train, "X_test": X_test, "y_test": y_test}
+        meta = create_dataset_meta("persist-test", spiral_params, X, y, len(X_train), len(X_test))
+
+        store1.save("persist-test", meta, arrays)
+
+        store2 = LocalFSDatasetStore(temp_storage_dir)
+
+        assert store2.exists("persist-test")
+        retrieved = store2.get_meta("persist-test")
+        assert retrieved is not None
+        assert retrieved.dataset_id == "persist-test"
+
+        datasets = store2.list_datasets()
+        assert "persist-test" in datasets
+
+
+@pytest.mark.integration
+class TestDatasetLifecycle:
+    """Tests for complete dataset lifecycle operations."""
+
+    def test_create_update_delete_lifecycle(self, temp_storage_dir: Path, spiral_params: SpiralParams):
+        """Test full create → verify → delete lifecycle."""
+        store = LocalFSDatasetStore(temp_storage_dir)
+
+        data = SpiralGenerator.generate(spiral_params)
+        X, y = data["X_full"], data["y_full"]
+        X_train, y_train = data["X_train"], data["y_train"]
+        X_test, y_test = data["X_test"], data["y_test"]
+
+        arrays = {"X_train": X_train, "y_train": y_train, "X_test": X_test, "y_test": y_test}
+        meta = create_dataset_meta("lifecycle-test", spiral_params, X, y, len(X_train), len(X_test))
+
+        store.save("lifecycle-test", meta, arrays)
+        assert store.exists("lifecycle-test")
+        assert "lifecycle-test" in store.list_datasets()
+
+        deleted = store.delete("lifecycle-test")
+        assert deleted
+        assert not store.exists("lifecycle-test")
+        assert "lifecycle-test" not in store.list_datasets()
+        assert store.get_meta("lifecycle-test") is None
+        assert store.get_artifact_bytes("lifecycle-test") is None
+
+    def test_multiple_datasets(self, temp_storage_dir: Path):
+        """Store and manage multiple datasets."""
+        store = LocalFSDatasetStore(temp_storage_dir)
+
+        for i, n_points in enumerate([50, 100, 200]):
+            params = SpiralParams(n_spirals=2, n_points_per_spiral=n_points, seed=i)
+            data = SpiralGenerator.generate(params)
+            X, y = data["X_full"], data["y_full"]
+            X_train, y_train = data["X_train"], data["y_train"]
+            X_test, y_test = data["X_test"], data["y_test"]
+
+            arrays = {"X_train": X_train, "y_train": y_train, "X_test": X_test, "y_test": y_test}
+            meta = create_dataset_meta(f"multi-{i}", params, X, y, len(X_train), len(X_test))
+
+            store.save(f"multi-{i}", meta, arrays)
+
+        datasets = store.list_datasets()
+        assert len(datasets) == 3
+        assert all(f"multi-{i}" in datasets for i in range(3))
+
+        store.delete("multi-1")
+        datasets = store.list_datasets()
+        assert len(datasets) == 2
+        assert "multi-1" not in datasets
+
+
+@pytest.mark.integration
+class TestChecksumVerification:
+    """Tests for checksum verification across storage operations."""
+
+    def test_checksum_consistency(self, temp_storage_dir: Path, spiral_params: SpiralParams):
+        """Checksum remains consistent through store/retrieve cycle."""
+        store = LocalFSDatasetStore(temp_storage_dir)
+
+        data = SpiralGenerator.generate(spiral_params)
+        X, y = data["X_full"], data["y_full"]
+        X_train, y_train = data["X_train"], data["y_train"]
+        X_test, y_test = data["X_test"], data["y_test"]
+
+        arrays = {"X_train": X_train, "y_train": y_train, "X_test": X_test, "y_test": y_test}
+
+        original_checksum = compute_checksum(arrays)
+
+        meta = create_dataset_meta("checksum-test", spiral_params, X, y, len(X_train), len(X_test))
+        meta.checksum = original_checksum
+
+        store.save("checksum-test", meta, arrays)
+
+        retrieved_meta = store.get_meta("checksum-test")
+        assert retrieved_meta is not None
+        assert retrieved_meta.checksum == original_checksum
+
+        artifact_bytes = store.get_artifact_bytes("checksum-test")
+        assert artifact_bytes is not None
+        loaded = np.load(io.BytesIO(artifact_bytes))
+        loaded_arrays = {k: loaded[k] for k in loaded.files}
+
+        loaded_checksum = compute_checksum(loaded_arrays)
+        assert loaded_checksum == original_checksum
+
+
+@pytest.mark.integration
+class TestReproducibility:
+    """Tests for dataset generation reproducibility."""
+
+    def test_seed_reproducibility(self):
+        """Same seed produces identical datasets."""
+        params1 = SpiralParams(n_spirals=2, n_points_per_spiral=100, seed=42)
+        params2 = SpiralParams(n_spirals=2, n_points_per_spiral=100, seed=42)
+
+        data1 = SpiralGenerator.generate(params1)
+        data2 = SpiralGenerator.generate(params2)
+
+        np.testing.assert_array_equal(data1["X_full"], data2["X_full"])
+        np.testing.assert_array_equal(data1["y_full"], data2["y_full"])
+
+    def test_different_seeds_produce_different_data(self):
+        """Different seeds produce different datasets."""
+        params1 = SpiralParams(n_spirals=2, n_points_per_spiral=100, seed=42)
+        params2 = SpiralParams(n_spirals=2, n_points_per_spiral=100, seed=43)
+
+        data1 = SpiralGenerator.generate(params1)
+        data2 = SpiralGenerator.generate(params2)
+
+        assert not np.allclose(data1["X_full"], data2["X_full"])

juniper_data/tests/performance/__init__.py

@@ -0,0 +1 @@
+"""Performance benchmark tests for juniper_data."""

juniper_data/tests/performance/test_generator_benchmarks.py

@@ -0,0 +1,178 @@
+#####################################################################################################################################################################################################
+# Project:       Juniper
+# Sub-Project:   JuniperData
+# Application:   juniper_data
+# File Name:     test_generator_benchmarks.py
+# Author:        Paul Calnon
+# Version:       0.4.2
+#
+# Date Created:  2026-02-25
+# Last Modified: 2026-02-25
+#
+# License:       MIT License
+# Copyright:     Copyright (c) 2024-2026 Paul Calnon
+#
+# Description:
+#    Performance benchmarks for dataset generators.
+#    Measures generation throughput (points/second) for each synthetic generator
+#    and scaling behavior across dataset sizes.
+#
+# Usage:
+#    # Run benchmarks with timing (disabled by default in addopts):
+#    pytest juniper_data/tests/performance/test_generator_benchmarks.py --benchmark-enable -v
+#
+#    # Run with autosave for regression tracking:
+#    pytest juniper_data/tests/performance/test_generator_benchmarks.py --benchmark-enable --benchmark-autosave
+#
+#    # Compare against saved baseline:
+#    pytest juniper_data/tests/performance/test_generator_benchmarks.py --benchmark-enable --benchmark-compare
+#
+# References:
+#    - RD-009: Performance Test Infrastructure
+#    - pytest-benchmark: https://pytest-benchmark.readthedocs.io/
+#####################################################################################################################################################################################################
+
+"""Performance benchmarks for dataset generators.
+
+Benchmarks measure generation throughput for each synthetic generator
+and scaling behavior with increasing dataset sizes. External-dependency
+generators (MNIST, ARC-AGI, CSV import) are excluded as they measure
+I/O rather than generation performance.
+"""
+
+import numpy as np
+import pytest
+
+from juniper_data.generators.checkerboard.generator import CheckerboardGenerator
+from juniper_data.generators.checkerboard.params import CheckerboardParams
+from juniper_data.generators.circles.generator import CirclesGenerator
+from juniper_data.generators.circles.params import CirclesParams
+from juniper_data.generators.gaussian.generator import GaussianGenerator
+from juniper_data.generators.gaussian.params import GaussianParams
+from juniper_data.generators.spiral.generator import SpiralGenerator
+from juniper_data.generators.spiral.params import SpiralParams
+from juniper_data.generators.xor.generator import XorGenerator
+from juniper_data.generators.xor.params import XorParams
+
+# ═══════════════════════════════════════════════════════════════════════════════
+# Generator Throughput Benchmarks
+# ═══════════════════════════════════════════════════════════════════════════════
+
+
+@pytest.mark.performance
+class TestGeneratorThroughput:
+    """Benchmark each synthetic generator with standard parameters.
+
+    Each test generates a dataset of ~1000 total points and validates
+    the output structure. The benchmark fixture handles timing and
+    iteration count automatically.
+    """
+
+    def test_spiral_generator(self, benchmark):
+        """Benchmark spiral dataset generation (1000 points)."""
+        params = SpiralParams(n_spirals=2, n_points_per_spiral=500, seed=42)
+        result = benchmark(SpiralGenerator.generate, params)
+        assert "X_train" in result
+        assert result["X_full"].shape[0] == 1000
+        assert result["X_train"].dtype == np.float32
+
+    def test_xor_generator(self, benchmark):
+        """Benchmark XOR dataset generation (1000 points)."""
+        params = XorParams(n_points_per_quadrant=250, seed=42)
+        result = benchmark(XorGenerator.generate, params)
+        assert "X_train" in result
+        assert result["X_full"].shape[0] == 1000
+        assert result["X_train"].dtype == np.float32
+
+    def test_gaussian_generator(self, benchmark):
+        """Benchmark Gaussian blobs dataset generation (1000 points)."""
+        params = GaussianParams(n_classes=2, n_samples_per_class=500, seed=42)
+        result = benchmark(GaussianGenerator.generate, params)
+        assert "X_train" in result
+        assert result["X_full"].shape[0] == 1000
+        assert result["X_train"].dtype == np.float32
+
+    def test_circles_generator(self, benchmark):
+        """Benchmark concentric circles dataset generation (1000 points)."""
+        params = CirclesParams(n_samples=1000, seed=42)
+        result = benchmark(CirclesGenerator.generate, params)
+        assert "X_train" in result
+        assert result["X_full"].shape[0] == 1000
+        assert result["X_train"].dtype == np.float32
+
+    def test_checkerboard_generator(self, benchmark):
+        """Benchmark checkerboard dataset generation (1000 points)."""
+        params = CheckerboardParams(n_samples=1000, seed=42)
+        result = benchmark(CheckerboardGenerator.generate, params)
+        assert "X_train" in result
+        assert result["X_full"].shape[0] == 1000
+        assert result["X_train"].dtype == np.float32
+
+
+# ═══════════════════════════════════════════════════════════════════════════════
+# Scaling Benchmarks
+# ═══════════════════════════════════════════════════════════════════════════════
+
+
+@pytest.mark.performance
+class TestGeneratorScaling:
+    """Benchmark generation throughput across dataset sizes.
+
+    Tests the spiral generator (representative of numpy-based generators)
+    with increasing point counts to characterize scaling behavior.
+    """
+
+    @pytest.mark.parametrize(
+        "n_points_per_spiral",
+        [100, 500, 1000, 5000],
+        ids=["200pts", "1000pts", "2000pts", "10000pts"],
+    )
+    def test_spiral_scaling(self, benchmark, n_points_per_spiral):
+        """Benchmark spiral generation at various dataset sizes."""
+        params = SpiralParams(n_spirals=2, n_points_per_spiral=n_points_per_spiral, seed=42)
+        result = benchmark(SpiralGenerator.generate, params)
+        assert result["X_full"].shape[0] == n_points_per_spiral * 2
+
+    @pytest.mark.parametrize(
+        "n_samples",
+        [100, 500, 1000, 5000],
+        ids=["100pts", "500pts", "1000pts", "5000pts"],
+    )
+    def test_gaussian_scaling(self, benchmark, n_samples):
+        """Benchmark Gaussian generation at various dataset sizes."""
+        params = GaussianParams(n_classes=2, n_samples_per_class=n_samples // 2, seed=42)
+        result = benchmark(GaussianGenerator.generate, params)
+        assert result["X_full"].shape[0] == n_samples
+
+
+# ═══════════════════════════════════════════════════════════════════════════════
+# Multi-Class Benchmarks
+# ═══════════════════════════════════════════════════════════════════════════════
+
+
+@pytest.mark.performance
+class TestMultiClassScaling:
+    """Benchmark generation throughput as class count increases."""
+
+    @pytest.mark.parametrize(
+        "n_spirals",
+        [2, 3, 5, 8],
+        ids=["2class", "3class", "5class", "8class"],
+    )
+    def test_spiral_class_scaling(self, benchmark, n_spirals):
+        """Benchmark spiral generation with varying class counts."""
+        params = SpiralParams(n_spirals=n_spirals, n_points_per_spiral=200, seed=42)
+        result = benchmark(SpiralGenerator.generate, params)
+        assert result["X_full"].shape[0] == n_spirals * 200
+        assert result["y_full"].shape[1] == n_spirals
+
+    @pytest.mark.parametrize(
+        "n_classes",
+        [2, 3, 5, 8],
+        ids=["2class", "3class", "5class", "8class"],
+    )
+    def test_gaussian_class_scaling(self, benchmark, n_classes):
+        """Benchmark Gaussian generation with varying class counts."""
+        params = GaussianParams(n_classes=n_classes, n_samples_per_class=100, seed=42)
+        result = benchmark(GaussianGenerator.generate, params)
+        assert result["X_full"].shape[0] == n_classes * 100