From 71ec78c7a3d1a00e8b98a0e0a2d48066358eb49d Mon Sep 17 00:00:00 2001
From: Eivind Jahren <ejah@equinor.com>
Date: Tue, 19 Nov 2024 15:06:17 +0100
Subject: [PATCH] Fix issue where dark_storage did not handle empty responses

---
 src/ert/dark_storage/common.py                |  3 +-
 .../unit_tests/dark_storage/test_common.py    | 32 ++++++++++++++++++-
 2 files changed, 33 insertions(+), 2 deletions(-)

diff --git a/src/ert/dark_storage/common.py b/src/ert/dark_storage/common.py
index c6871c3e452..3ef7ac3f1ea 100644
--- a/src/ert/dark_storage/common.py
+++ b/src/ert/dark_storage/common.py
@@ -6,6 +6,7 @@
 import pandas as pd
 import polars
 import xarray as xr
+from polars.exceptions import ColumnNotFoundError
 
 from ert.config import GenDataConfig, GenKwConfig
 from ert.config.field import Field
@@ -184,7 +185,7 @@ def data_for_key(
                 return data.astype(float)
             except ValueError:
                 return data
-        except (ValueError, KeyError):
+        except (ValueError, KeyError, ColumnNotFoundError):
             return pd.DataFrame()
 
     return pd.DataFrame()
diff --git a/tests/ert/unit_tests/dark_storage/test_common.py b/tests/ert/unit_tests/dark_storage/test_common.py
index b05ed558c79..1473e3e80e1 100644
--- a/tests/ert/unit_tests/dark_storage/test_common.py
+++ b/tests/ert/unit_tests/dark_storage/test_common.py
@@ -1,7 +1,8 @@
 import pandas as pd
+import polars
 import pytest
 
-from ert.config import SummaryConfig
+from ert.config import GenDataConfig, SummaryConfig
 from ert.dark_storage.common import data_for_key
 from ert.storage import open_storage
 from tests.ert.unit_tests.config.summary_generator import (
@@ -70,3 +71,32 @@ def test_data_for_key_gives_mean_for_duplicate_values(tmp_path):
         assert df[pd.Timestamp("2014-01-16 05:00:00")][0] == pytest.approx(
             (value1 + value2) / 2
         )
+
+
+def test_data_for_key_returns_empty_gen_data_config(tmp_path):
+    with open_storage(tmp_path / "storage", mode="w") as storage:
+        gen_data_config = GenDataConfig(keys=["response"])
+        experiment = storage.create_experiment(
+            observations={},
+            parameters=[],
+            responses=[gen_data_config],
+        )
+        ensemble = experiment.create_ensemble(name="ensemble", ensemble_size=1)
+
+        data = data_for_key(ensemble, "response@0")
+        assert data.empty
+
+        ensemble.save_response(
+            "gen_data",
+            polars.DataFrame(
+                {
+                    "response_key": "response",
+                    "report_step": polars.Series([0], dtype=polars.UInt16),
+                    "index": polars.Series([0], dtype=polars.UInt16),
+                    "values": polars.Series([0.0], dtype=polars.Float32),
+                }
+            ),
+            0,
+        )
+        data = data_for_key(ensemble, "response@0")
+        assert not data.empty