Begin work on tests

Author: BrianMichell

tests/integration/test_segy_import_export_masked.py

@@ -282,11 +282,16 @@ def generate_selection_mask(selection_conf: SelectionMaskConfig, grid_conf: Grid
 
 
 @pytest.fixture
-def export_masked_path(tmp_path_factory: pytest.TempPathFactory) -> Path:
+def export_masked_path(tmp_path_factory: pytest.TempPathFactory, raw_headers_env: None) -> Path:  # noqa: ARG001
     """Fixture that generates temp directory for export tests."""
+    # Create path suffix based on current raw headers environment variable
+    # raw_headers_env dependency ensures the environment variable is set before this runs
+    raw_headers_enabled = os.getenv("MDIO__DO_RAW_HEADERS") == "1"
+    path_suffix = "with_raw_headers" if raw_headers_enabled else "without_raw_headers"
+    
     if DEBUG_MODE:
-        return Path("TMP/export_masked")
-    return tmp_path_factory.getbasetemp() / "export_masked"
+        return Path(f"TMP/export_masked_{path_suffix}")
+    return tmp_path_factory.getbasetemp() / f"export_masked_{path_suffix}"
 
 
 @pytest.fixture
@@ -300,8 +305,33 @@ def raw_headers_env(request: pytest.FixtureRequest) -> None:
 
     yield
 
-    # Cleanup after test
+    # Cleanup after test - both environment variable and template state
     os.environ.pop("MDIO__DO_RAW_HEADERS", None)
+    
+    # Clean up any template modifications to ensure test isolation
+    from mdio.builder.template_registry import TemplateRegistry
+    registry = TemplateRegistry.get_instance()
+    
+    # Reset any templates that might have been modified with raw headers
+    template_names = ["PostStack2DTime", "PostStack3DTime", "PreStackCdpOffsetGathers2DTime", 
+                     "PreStackCdpOffsetGathers3DTime", "PreStackShotGathers2DTime", 
+                     "PreStackShotGathers3DTime", "PreStackCocaGathers3DTime"]
+    
+    for template_name in template_names:
+        try:
+            template = registry.get(template_name)
+            # Remove raw headers enhancement if present
+            if hasattr(template, "_mdio_raw_headers_enhanced"):
+                delattr(template, "_mdio_raw_headers_enhanced")
+                # The enhancement is applied by monkey-patching _add_variables
+                # We need to restore it to the original method from the class
+                # Since we can't easily restore the exact original, we'll get a fresh instance
+                template_class = type(template)
+                if hasattr(template_class, '_add_variables'):
+                    template._add_variables = template_class._add_variables.__get__(template, template_class)
+        except KeyError:
+            # Template not found, skip
+            continue
 
 
 @pytest.mark.parametrize(
@@ -471,3 +501,64 @@ def test_export_masked(
         # https://github.com/TGSAI/mdio-python/issues/610
         assert_array_equal(actual_sgy.trace[:].header, expected_sgy.trace[expected_trc_idx].header)
         assert_array_equal(actual_sgy.trace[:].sample, expected_sgy.trace[expected_trc_idx].sample)
+
+    def test_raw_headers_byte_preservation(
+        self,
+        test_conf: MaskedExportConfig,
+        export_masked_path: Path,
+        raw_headers_env: None,  # noqa: ARG002
+    ) -> None:
+        """Test that raw headers are preserved byte-for-byte when MDIO__DO_RAW_HEADERS=1."""
+        grid_conf, segy_factory_conf, _, _ = test_conf
+        segy_path = export_masked_path / f"{grid_conf.name}.sgy"
+        mdio_path = export_masked_path / f"{grid_conf.name}.mdio"
+
+        # Open MDIO dataset
+        ds = open_mdio(mdio_path)
+
+        # Check if raw_headers should exist based on environment variable
+        has_raw_headers = "raw_headers" in ds.data_vars
+        if os.getenv("MDIO__DO_RAW_HEADERS") == "1":
+            assert has_raw_headers, "raw_headers should be present when MDIO__DO_RAW_HEADERS=1"
+        else:
+            assert not has_raw_headers, f"raw_headers should not be present when MDIO__DO_RAW_HEADERS is not set\n {ds}"
+            return  # Exit early if raw_headers are not expected
+        
+        # Get data (only if raw_headers exist)
+        raw_headers_data = ds.raw_headers.values
+        trace_mask = ds.trace_mask.values
+        
+        # Verify 240-byte headers
+        assert raw_headers_data.dtype.itemsize == 240, f"Expected 240-byte headers, got {raw_headers_data.dtype.itemsize}"
+        
+        # Read raw bytes directly from SEG-Y file
+        def read_segy_trace_header(trace_index: int) -> bytes:
+            """Read 240-byte trace header directly from SEG-Y file."""
+            with open(segy_path, 'rb') as f:
+                # Skip text (3200) + binary (400) headers = 3600 bytes
+                f.seek(3600)
+                # Each trace: 240 byte header + (num_samples * 4) byte samples
+                trace_size = 240 + (segy_factory_conf.num_samples * 4)
+                trace_offset = trace_index * trace_size
+                f.seek(trace_offset, 1)  # Seek relative to current position
+                return f.read(240)
+        
+        # Compare all valid traces byte-by-byte
+        segy_trace_idx = 0
+        flat_mask = trace_mask.ravel()
+        flat_raw_headers = raw_headers_data.reshape(-1, raw_headers_data.shape[-1])
+        
+        for grid_idx in range(flat_mask.size):
+            if not flat_mask[grid_idx]:
+                continue
+                
+            # Get MDIO header as bytes
+            mdio_header_bytes = np.frombuffer(flat_raw_headers[grid_idx].tobytes(), dtype=np.uint8)
+            
+            # Get SEG-Y header as raw bytes directly from file
+            segy_raw_header_bytes = read_segy_trace_header(segy_trace_idx)
+            segy_header_bytes = np.frombuffer(segy_raw_header_bytes, dtype=np.uint8)
+            
+            # Compare byte-by-byte
+            assert_array_equal(mdio_header_bytes, segy_header_bytes)
+            segy_trace_idx += 1