From a414e6efc0d1558c3bfea18f87886f7fc5d90bab Mon Sep 17 00:00:00 2001
From: sebhmg <sebastienh@mirageoscience.com>
Date: Thu, 23 Nov 2023 12:34:37 -0500
Subject: [PATCH] add back the compression option, plus tests

---
 omf/scripts/omf_to_geoh5.py       | 23 ++++++++--
 tests/script_omf_to_geoh5_test.py | 74 ++++++++++++++++++++++++++++---
 2 files changed, 89 insertions(+), 8 deletions(-)

diff --git a/omf/scripts/omf_to_geoh5.py b/omf/scripts/omf_to_geoh5.py
index e76e5a54..b7ecc7e1 100644
--- a/omf/scripts/omf_to_geoh5.py
+++ b/omf/scripts/omf_to_geoh5.py
@@ -14,8 +14,25 @@ def run():
         prog="omf_to_geoh5",
         description="Converts an OMF file to a new geoh5 file.",
     )
-    parser.add_argument("omf_file", type=Path)
-    parser.add_argument("-o", "--out", type=Path, required=False, default=None)
+    parser.add_argument("omf_file", type=Path, help="Path to the OMF file to convert.")
+    parser.add_argument(
+        "-o",
+        "--out",
+        type=Path,
+        required=False,
+        default=None,
+        help=(
+            "Path to the output geoh5 file. If not specified, create the output file "
+            "at the same location as the input file, but with the geoh5 extension."
+        ),
+    )
+    parser.add_argument(
+        "--gzip",
+        type=int,
+        choices=range(0, 10),
+        default=5,
+        help="Gzip compression level (0-9) for h5 data.",
+    )
     args = parser.parse_args()
 
     omf_filepath = args.omf_file
@@ -32,7 +49,7 @@ def run():
         sys.exit(1)
 
     reader = OMFReader(str(omf_filepath.absolute()))
-    GeoH5Writer(reader.get_project(), output_filepath, compression=compression)
+    GeoH5Writer(reader.get_project(), output_filepath, compression=args.gzip)
     _logger.info("geoh5 file created: %s", output_filepath)
 
 
diff --git a/tests/script_omf_to_geoh5_test.py b/tests/script_omf_to_geoh5_test.py
index b4f4ec56..aa3ba026 100644
--- a/tests/script_omf_to_geoh5_test.py
+++ b/tests/script_omf_to_geoh5_test.py
@@ -12,6 +12,13 @@
     scope="module", name="omf_input_path", params=["test file.omf", "test_file.omf"]
 )
 def omf_input_path_fixture(request, tmp_path_factory) -> Path:
+    omf_path = tmp_path_factory.mktemp("input") / request.param
+    create_omf_file(omf_path)
+    return omf_path
+
+
+def create_omf_file(omf_file_path: Path) -> None:
+    """Create an OMF file with random data."""
     points = omf.PointSetElement(
         name="Random Points",
         description="Just random points",
@@ -23,9 +30,8 @@ def omf_input_path_fixture(request, tmp_path_factory) -> Path:
         ],
     )
 
-    file_path = tmp_path_factory.mktemp("input") / request.param
-    omf.OMFWriter(points, str(file_path))
-    return file_path
+    omf.OMFWriter(points, str(omf_file_path))
+    assert omf_file_path.exists()
 
 
 def test_omf_to_geoh5_without_output_name(omf_input_path: Path):
@@ -41,7 +47,7 @@ def test_omf_to_geoh5_without_output_name(omf_input_path: Path):
     "output_name", ["my_output.geoh5", "my output.geoh5", "my_output", "my output"]
 )
 def test_omf_to_geoh5_with_output_name(
-    tmp_path, monkeypatch, omf_input_path: Path, output_name: str
+    tmp_path: Path, monkeypatch, omf_input_path: Path, output_name: str
 ):
     """Test the omf_to_geoh5 script."""
 
@@ -63,7 +69,7 @@ def test_omf_to_geoh5_with_output_name(
     "output_name", ["my_output.geoh5", "my output.geoh5", "my_output", "my output"]
 )
 def test_omf_to_geoh5_with_absolute_output_path(
-    tmp_path, omf_input_path: Path, output_name: str
+    tmp_path: Path, omf_input_path: Path, output_name: str
 ):
     """Test the omf_to_geoh5 script."""
 
@@ -84,3 +90,61 @@ def test_omf_to_geoh5_with_absolute_output_path(
     if not expected_output.suffix:
         expected_output = expected_output.with_suffix(".geoh5")
     assert expected_output.exists()
+
+
+@pytest.mark.parametrize("gzip_level", range(0, 10))
+def test_omf_to_geoh5_with_gzip_level(tmp_path: Path, gzip_level: int):
+    """Test the omf_to_geoh5 script."""
+
+    omf_path = tmp_path / "test_file.omf"
+    create_omf_file(omf_path)
+    output_name = f"{omf_path.stem}_{gzip_level}.geoh5"
+    output_dir = tmp_path / "output"
+    output_path = output_dir / output_name
+    output_dir.mkdir()
+    with patch(
+        "sys.argv",
+        [
+            "omf_to_geoh5",
+            str(omf_path),
+            "--gzip",
+            f"{gzip_level}",
+            "-o",
+            f"{output_path.absolute()}",
+        ],
+    ):
+        omf_to_geoh5.run()
+
+    assert output_path.exists()
+
+
+def test_omf_to_geoh5_with_gzip_level_too_high(capsys, tmp_path: Path):
+    """Test the omf_to_geoh5 script."""
+
+    omf_path = tmp_path / "test_file.omf"
+    create_omf_file(omf_path)
+    output_name = omf_path.with_suffix(".geoh5").name
+    output_dir = tmp_path / "output"
+    output_path = output_dir / output_name
+    output_dir.mkdir()
+    with pytest.raises(SystemExit) as captured_exception:
+        with patch(
+            "sys.argv",
+            [
+                "omf_to_geoh5",
+                str(omf_path),
+                "--gzip",
+                "10",
+                "-o",
+                f"{output_path.absolute()}",
+            ],
+        ):
+            omf_to_geoh5.run()
+
+    assert not output_path.exists()
+    assert captured_exception.value.code == 2
+    captured_err = capsys.readouterr().err
+    assert any(
+        "error: argument --gzip: invalid choice: 10" in line
+        for line in captured_err.splitlines()
+    )