Begin adding tests

BrianMichell · BrianMichell · commit 73bf9362abf7 · 2025-11-17T21:03:58.000Z
diff --git a/tests/integration/transpose_writers/test_chunking.py b/tests/integration/transpose_writers/test_chunking.py
@@ -0,0 +1,145 @@
+"""Integration tests for chunking module."""
+
+from __future__ import annotations
+
+from typing import TYPE_CHECKING
+
+import dask
+import numpy as np
+import pytest
+from segy.factory import SegyFactory
+from segy.schema import HeaderField
+from segy.schema import ScalarType
+from segy.standards import get_segy_standard
+
+from mdio.api.io import open_mdio
+from mdio.builder.schemas.chunk_grid import RegularChunkGrid
+from mdio.builder.schemas.chunk_grid import RegularChunkShape
+from mdio.builder.schemas.compressors import Blosc
+from mdio.builder.template_registry import TemplateRegistry
+from mdio.converters.segy import segy_to_mdio
+from mdio.transpose_writers.chunking import from_variable
+
+if TYPE_CHECKING:
+    from pathlib import Path
+
+    from segy.schema import SegySpec
+
+
+dask.config.set(scheduler="synchronous")
+
+
+@pytest.fixture
+def synthetic_segy_spec() -> SegySpec:
+    """SEG-Y specification for synthetic testing."""
+    trace_header_fields = [
+        HeaderField(name="inline", byte=189, format=ScalarType.INT32),
+        HeaderField(name="crossline", byte=193, format=ScalarType.INT32),
+        HeaderField(name="cdp_x", byte=181, format=ScalarType.INT32),
+        HeaderField(name="cdp_y", byte=185, format=ScalarType.INT32),
+    ]
+    return get_segy_standard(1.0).customize(trace_header_fields=trace_header_fields)
+
+
+@pytest.fixture
+def synthetic_segy_file(fake_segy_tmp: Path, synthetic_segy_spec: SegySpec) -> Path:
+    """Create a small synthetic 3D SEG-Y file."""
+    segy_path = fake_segy_tmp / "synthetic_3d.sgy"
+    inlines, crosslines, num_samples = np.arange(1, 9), np.arange(1, 17), 64
+    num_traces = len(inlines) * len(crosslines)
+
+    factory = SegyFactory(spec=synthetic_segy_spec, sample_interval=4000, samples_per_trace=num_samples)
+    headers = factory.create_trace_header_template(num_traces)
+    samples = factory.create_trace_sample_template(num_traces)
+
+    trace_idx = 0
+    for inline in inlines:
+        for crossline in crosslines:
+            headers["inline"][trace_idx] = inline
+            headers["crossline"][trace_idx] = crossline
+            headers["cdp_x"][trace_idx] = inline * 100
+            headers["cdp_y"][trace_idx] = crossline * 100
+            headers["coordinate_scalar"][trace_idx] = -100
+            samples[trace_idx] = np.sin(np.linspace(0, 2 * np.pi * inline, num_samples)) * crossline
+            trace_idx += 1
+
+    with segy_path.open(mode="wb") as fp:
+        fp.write(factory.create_textual_header())
+        fp.write(factory.create_binary_header())
+        fp.write(factory.create_traces(headers, samples))
+
+    return segy_path
+
+
+@pytest.fixture
+def mdio_dataset(synthetic_segy_file: Path, synthetic_segy_spec: SegySpec, zarr_tmp: Path) -> Path:
+    """Convert synthetic SEG-Y to MDIO."""
+    mdio_path = zarr_tmp / "test_dataset.mdio"
+    segy_to_mdio(
+        segy_spec=synthetic_segy_spec,
+        mdio_template=TemplateRegistry().get("PostStack3DTime"),
+        input_path=synthetic_segy_file,
+        output_path=mdio_path,
+        overwrite=True,
+    )
+    return mdio_path
+
+
+def test_single_variable_rechunk(mdio_dataset: Path) -> None:
+    """Test creating a single rechunked variable."""
+    chunk_grid = RegularChunkGrid(configuration=RegularChunkShape(chunk_shape=(2, 16, 64)))
+    compressor = Blosc(cname="zstd", clevel=5, shuffle="shuffle")
+
+    from_variable(
+        dataset_path=mdio_dataset,
+        source_variable="amplitude",
+        new_variable="fast_inline",
+        chunk_grid=chunk_grid,
+        compressor=compressor,
+        copy_metadata=True,
+    )
+
+    ds = open_mdio(mdio_dataset)
+
+    # Verify new variable exists with correct chunks
+    assert "fast_inline" in ds.data_vars
+    assert ds["fast_inline"].encoding["chunks"] == (2, 16, 64)
+
+    # Verify data integrity
+    np.testing.assert_array_equal(ds["amplitude"].values, ds["fast_inline"].values)
+
+    # Verify metadata copied
+    assert ds["fast_inline"].attrs == ds["amplitude"].attrs
+
+
+def test_multiple_variables_with_broadcasting(mdio_dataset: Path) -> None:
+    """Test creating multiple variables with different settings and broadcasting."""
+    new_variables = ["fast_inline", "fast_crossline", "fast_time"]
+    chunk_grids = [
+        RegularChunkGrid(configuration=RegularChunkShape(chunk_shape=(2, 16, 64))),
+        RegularChunkGrid(configuration=RegularChunkShape(chunk_shape=(8, 4, 64))),
+        RegularChunkGrid(configuration=RegularChunkShape(chunk_shape=(8, 16, 16))),
+    ]
+    compressor = Blosc(cname="lz4", clevel=3, shuffle="shuffle")  # Single compressor broadcasts
+
+    from_variable(
+        dataset_path=mdio_dataset,
+        source_variable="amplitude",
+        new_variable=new_variables,
+        chunk_grid=chunk_grids,
+        compressor=compressor,
+        copy_metadata=False,
+    )
+
+    ds = open_mdio(mdio_dataset)
+
+    # Verify all variables created with correct chunks
+    expected_chunks = [(2, 16, 64), (8, 4, 64), (8, 16, 16)]
+    for var_name, chunks in zip(new_variables, expected_chunks, strict=True):
+        assert var_name in ds.data_vars
+        assert ds[var_name].encoding["chunks"] == chunks
+        assert "compressors" in ds[var_name].encoding
+        np.testing.assert_array_equal(ds["amplitude"].values, ds[var_name].values)
+
+        # Metadata should not be copied
+        assert len(ds[var_name].attrs) == 0 or ds[var_name].attrs != ds["amplitude"].attrs
diff --git a/tests/unit/transpose_writers/test_chunking.py b/tests/unit/transpose_writers/test_chunking.py
@@ -0,0 +1,153 @@
+"""Unit tests for chunking module helper functions."""
+
+from __future__ import annotations
+
+from typing import TYPE_CHECKING
+
+import numpy as np
+import pytest
+import xarray as xr
+import zarr
+
+from mdio.api.io import open_mdio
+from mdio.api.io import to_mdio
+from mdio.builder.schemas.chunk_grid import RegularChunkGrid
+from mdio.builder.schemas.chunk_grid import RegularChunkShape
+from mdio.builder.schemas.compressors import Blosc
+from mdio.constants import ZarrFormat
+from mdio.transpose_writers.chunking import _normalize_chunk_grid
+from mdio.transpose_writers.chunking import _normalize_compressor
+from mdio.transpose_writers.chunking import _normalize_new_variable
+from mdio.transpose_writers.chunking import _remove_fillvalue_attrs
+from mdio.transpose_writers.chunking import _validate_inputs
+
+if TYPE_CHECKING:
+    from pathlib import Path
+
+
+class TestRemoveFillvalueAttrs:
+    """Tests for _remove_fillvalue_attrs helper function."""
+
+    @pytest.mark.parametrize("zarr_format", [ZarrFormat.V2, ZarrFormat.V3])
+    def test_remove_fillvalue_after_mdio_serialization(self, tmp_path: Path, zarr_format: ZarrFormat) -> None:
+        """Test that _FillValue is removed after MDIO serialization in both Zarr v2 and v3."""
+        # Create dataset with NaN values (will add _FillValue on serialization)
+        data = np.array([[1.0, 2.0, np.nan], [3.0, np.nan, 4.0]], dtype=np.float32)
+        ds = xr.Dataset(
+            {"var1": (["x", "y"], data, {"units": "meters"})},
+            coords={"x": (["x"], [0, 1], {"axis": "X"})},
+        )
+
+        # Write and read back with MDIO
+        mdio_path = tmp_path / f"test_{zarr_format}.mdio"
+        with zarr.config.set(default_zarr_format=zarr_format):
+            to_mdio(ds, mdio_path, mode="w")
+            ds_read = open_mdio(mdio_path)
+
+            # Apply function and verify _FillValue is removed everywhere
+            _remove_fillvalue_attrs(ds_read)
+
+            for var_name in list(ds_read.data_vars) + list(ds_read.coords):
+                assert "_FillValue" not in ds_read[var_name].attrs
+
+            # Verify other attributes preserved
+            assert ds_read["var1"].attrs["units"] == "meters"
+            assert ds_read["x"].attrs["axis"] == "X"
+
+
+class TestValidateInputs:
+    """Tests for _validate_inputs helper function."""
+
+    @pytest.mark.parametrize(
+        ("new_variable", "chunk_grid", "compressor", "should_pass"),
+        [
+            ("var1", "grid", "comp", True),
+            (["var1", "var2"], "grid", "comp", True),
+            ("var1", "grid", None, True),
+            (123, "grid", "comp", False),
+            ([], "grid", "comp", False),
+            ("var1", [], "comp", False),
+            ("var1", "grid", [], False),
+        ],
+    )
+    def test_validation(
+        self,
+        new_variable: str | list | int,
+        chunk_grid: str | list,
+        compressor: str | list | None,
+        should_pass: bool,
+    ) -> None:
+        """Test input validation with various combinations."""
+        grid = RegularChunkGrid(configuration=RegularChunkShape(chunk_shape=(10, 10)))
+        comp = Blosc(cname="zstd", clevel=5, shuffle="shuffle")
+
+        # Replace placeholders
+        if chunk_grid == "grid":
+            chunk_grid = grid
+        if compressor == "comp":
+            compressor = comp
+
+        if should_pass:
+            _validate_inputs(new_variable, chunk_grid, compressor)  # type: ignore[arg-type]
+        else:
+            with pytest.raises((TypeError, ValueError)):
+                _validate_inputs(new_variable, chunk_grid, compressor)  # type: ignore[arg-type]
+
+
+class TestNormalizeNewVariable:
+    """Tests for _normalize_new_variable helper function."""
+
+    @pytest.mark.parametrize(
+        ("input_value", "expected"),
+        [
+            ("var1", ["var1"]),
+            (["var1", "var2"], ["var1", "var2"]),
+        ],
+    )
+    def test_normalize(self, input_value: str | list[str], expected: list[str]) -> None:
+        """Test new_variable normalization."""
+        result = _normalize_new_variable(input_value)
+        assert result == expected
+
+
+class TestNormalizeChunkGrid:
+    """Tests for _normalize_chunk_grid helper function."""
+
+    def test_broadcast_and_match(self) -> None:
+        """Test chunk grid broadcasting and matching."""
+        grid = RegularChunkGrid(configuration=RegularChunkShape(chunk_shape=(10, 10, 10)))
+        grids = [grid] * 3
+
+        # Single grid broadcasts
+        assert len(_normalize_chunk_grid(grid, 3)) == 3
+
+        # List matches length
+        result = _normalize_chunk_grid(grids, 3)
+        assert len(result) == 3
+        assert result == grids
+
+        # Mismatch raises error
+        with pytest.raises(ValueError):
+            _normalize_chunk_grid([grid, grid], 3)
+
+
+class TestNormalizeCompressor:
+    """Tests for _normalize_compressor helper function."""
+
+    def test_broadcast_and_match(self) -> None:
+        """Test compressor broadcasting and matching."""
+        comp = Blosc(cname="zstd", clevel=5, shuffle="shuffle")
+
+        # None broadcasts
+        assert all(c is None for c in _normalize_compressor(None, 3))
+
+        # Single compressor broadcasts
+        assert len(_normalize_compressor(comp, 3)) == 3
+
+        # List with None entries
+        result = _normalize_compressor([comp, None, comp], 3)
+        assert result[1] is None
+
+        # Mismatch raises error
+        with pytest.raises(ValueError):
+            _normalize_compressor([comp, comp], 3)
