Add storage migration for removing is_improvement flag (#14117)

Author: eqbech

src/ert/storage/local_storage.py

@@ -31,7 +31,7 @@
 
 logger = logging.getLogger(__name__)
 
-_LOCAL_STORAGE_VERSION = 36
+_LOCAL_STORAGE_VERSION = 37
 
 
 def open_storage(
@@ -641,6 +641,7 @@ def _migrate(self, version: int) -> None:
             to34,
             to35,
             to36,
+            to37,
         )
 
         try:  # ruff: ignore[too-many-statements-in-try-clause]
@@ -706,6 +707,7 @@ def _migrate(self, version: int) -> None:
                     33: to34,
                     34: to35,
                     35: to36,
+                    36: to37,
                 }
                 for from_version in range(version, _LOCAL_STORAGE_VERSION):
                     migrations[from_version].migrate(self.path)

src/ert/storage/migration/to37.py

@@ -0,0 +1,27 @@
+from __future__ import annotations
+
+import json
+import logging
+from pathlib import Path
+
+logger = logging.getLogger(__name__)
+
+info = "Remove is_improvement flag from ensembles"
+
+
+def migrate(path: Path) -> None:
+    ensemble_dir = path / "ensembles"
+    if not ensemble_dir.exists():
+        return
+    for ens_dir in ensemble_dir.iterdir():
+        index_file = ens_dir / "index.json"
+        if not index_file.exists():
+            continue
+
+        index_data = json.loads(index_file.read_text(encoding="utf-8"))
+        ensemble_data = index_data.get("ensemble", {})
+
+        if "is_improvement" in ensemble_data:
+            ensemble_data.pop("is_improvement")
+            index_file.write_text(json.dumps(index_data, indent=2), encoding="utf-8")
+            logger.info("Removed is_improvement flag from %s", index_file)

tests/ert/unit_tests/storage/migration/test_to37.py

@@ -0,0 +1,79 @@
+import json
+from pathlib import Path
+
+import pytest
+
+from ert.storage.migration.to37 import migrate
+
+
+def _write_index(ens_path: Path, index_data: dict) -> Path:
+    ens_path.mkdir(parents=True)
+    index_file = ens_path / "index.json"
+    index_file.write_text(json.dumps(index_data), encoding="utf-8")
+    return index_file
+
+
+@pytest.mark.parametrize(
+    ("original", "expected"),
+    [
+        pytest.param(
+            {
+                "id": "ens-id",
+                "ensemble": {"name": "batch_0", "iteration": 0, "is_improvement": True},
+            },
+            {"id": "ens-id", "ensemble": {"name": "batch_0", "iteration": 0}},
+            id="removes_is_improvement",
+        ),
+        pytest.param(
+            {"ensemble": {"name": "batch_1", "iteration": 1, "is_improvement": False}},
+            {"ensemble": {"name": "batch_1", "iteration": 1}},
+            id="removes_is_improvement_false",
+        ),
+        pytest.param(
+            {"ensemble": {"name": "batch_5", "iteration": 3, "is_improvement": True}},
+            {"ensemble": {"name": "batch_5", "iteration": 3}},
+            id="removes_is_improvement_other_values",
+        ),
+        pytest.param(
+            {"ensemble": {"name": "batch_0", "iteration": 0}},
+            {"ensemble": {"name": "batch_0", "iteration": 0}},
+            id="leaves_untouched_when_missing",
+        ),
+    ],
+)
+def test_that_migration_updates_ensemble_index(tmp_path, original, expected):
+    root = tmp_path / "project"
+    index_file = _write_index(root / "ensembles" / "ensemble_1", original)
+
+    migrate(root)
+
+    assert json.loads(index_file.read_text(encoding="utf-8")) == expected
+
+
+def test_that_migration_does_not_fail_on_unexpectedly_structured_dirs(tmp_path):
+    root = tmp_path / "project"
+    root.mkdir()
+    migrate(root)
+
+    ensembles_dir = root / "ensembles"
+    ensembles_dir.mkdir()
+    migrate(root)
+
+    not_an_ensemble = ensembles_dir / "not_a_directory.json"
+    not_an_ensemble.write_text("{}", encoding="utf-8")
+    migrate(root)
+
+    ensemble_without_index_json = ensembles_dir / "ensemble_1"
+    ensemble_without_index_json.mkdir()
+    migrate(root)
+
+    assert not_an_ensemble.read_text(encoding="utf-8") == "{}"
+
+
+def test_that_migration_does_not_fail_on_index_without_ensemble_entry(tmp_path):
+    root = tmp_path / "project"
+    index_file = _write_index(root / "ensembles" / "ensemble_1", {"id": "ens-id"})
+
+    migrate(root)
+
+    assert json.loads(index_file.read_text(encoding="utf-8")) == {"id": "ens-id"}