Skip to content

Commit

Permalink
Make egg test use new storage
Browse files Browse the repository at this point in the history
  • Loading branch information
yngve-sk committed Jan 8, 2025
1 parent 6710ce7 commit b3761a4
Showing 1 changed file with 23 additions and 0 deletions.
23 changes: 23 additions & 0 deletions tests/everest/test_egg_simulation.py
Original file line number Diff line number Diff line change
@@ -1,6 +1,8 @@
import io
import json
import os

import polars
import pytest

from ert.config import ErtConfig, QueueSystem
Expand Down Expand Up @@ -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,
Expand Down

0 comments on commit b3761a4

Please sign in to comment.