From b3761a4d1c241ad9c12bc1a6c97be0a6796f6a18 Mon Sep 17 00:00:00 2001
From: "Yngve S. Kristiansen" <ynk@equinor.com>
Date: Wed, 8 Jan 2025 13:28:00 +0100
Subject: [PATCH] Make egg test use new storage

---
 tests/everest/test_egg_simulation.py | 23 +++++++++++++++++++++++
 1 file changed, 23 insertions(+)

diff --git a/tests/everest/test_egg_simulation.py b/tests/everest/test_egg_simulation.py
index 853b43bb8f4..6dc46690a69 100644
--- a/tests/everest/test_egg_simulation.py
+++ b/tests/everest/test_egg_simulation.py
@@ -1,6 +1,8 @@
+import io
 import json
 import os
 
+import polars
 import pytest
 
 from ert.config import ErtConfig, QueueSystem
@@ -804,6 +806,27 @@ def sweetcallbackofmine(self, *args, **kwargs):
 
     assert cbtracker.called
 
+    best_batch = [b for b in run_model.ever_storage.data.batches if b.is_improvement][
+        -1
+    ]
+
+    def _df_to_string(df: polars.DataFrame):
+        strbuf = io.StringIO()
+        schema = df.schema
+        df.with_columns(
+            polars.col(c) for c in df.columns if schema[c] == polars.Float32
+        ).write_csv(strbuf)
+
+        return strbuf.getvalue()
+
+    best_objectives_csv = _df_to_string(best_batch.perturbation_objectives)
+    best_objective_gradients_csv = _df_to_string(best_batch.batch_objective_gradient)
+    best_controls = _df_to_string(best_batch.realization_controls)
+
+    snapshot.assert_match(best_controls, "best_controls")
+    snapshot.assert_match(best_objectives_csv, "best_objectives_csv")
+    snapshot.assert_match(best_objective_gradients_csv, "best_objective_gradients_csv")
+
     # 2DO rewrite for everest storage
     # data = export_data(
     #    export_config=config.export,