Skip to content

Commit

Permalink
Small logic fix in data_for_key
Browse files Browse the repository at this point in the history
  • Loading branch information
yngve-sk committed Nov 28, 2024
1 parent 4dc583e commit 1aacc17
Show file tree
Hide file tree
Showing 2 changed files with 30 additions and 1 deletion.
3 changes: 2 additions & 1 deletion src/ert/dark_storage/common.py
Original file line number Diff line number Diff line change
Expand Up @@ -116,7 +116,8 @@ def data_for_key(
response_key = next((k for k in response_key_to_response_type if k == key), None)
if response_key is None:
response_key = next(
(k for k in response_key_to_response_type if k in key), None
(k for k in response_key_to_response_type if k in key and key != f"{k}H"),
None,
)

if response_key is not None:
Expand Down
28 changes: 28 additions & 0 deletions tests/ert/unit_tests/dark_storage/test_common.py
Original file line number Diff line number Diff line change
@@ -1,3 +1,5 @@
import datetime

import pandas as pd
import polars
import pytest
Expand Down Expand Up @@ -73,6 +75,32 @@ def test_data_for_key_gives_mean_for_duplicate_values(tmp_path):
)


def test_data_for_key_doesnt_mistake_history_for_response(tmp_path):
with open_storage(tmp_path / "storage", mode="w") as storage:
summary_config = SummaryConfig(
name="summary", input_files=["CASE"], keys=["FGPR"]
)
experiment = storage.create_experiment(responses=[summary_config])
ensemble = experiment.create_ensemble(name="ensemble", ensemble_size=1)
ensemble.save_response(
"summary",
polars.DataFrame(
{
"response_key": ["FGPR", "FGPR"],
"time": polars.Series(
[datetime.datetime(2000, 1, 1), datetime.datetime(2000, 1, 2)],
dtype=polars.Datetime("ms"),
),
"values": polars.Series([0.0, 1.0], dtype=polars.Float32),
}
),
0,
)
ensemble.refresh_ensemble_state()

assert data_for_key(ensemble, "FGPRH").empty


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"])
Expand Down

0 comments on commit 1aacc17

Please sign in to comment.