Skip to content

Commit

Permalink
Add ensemble plot test for correct rate draw style
Browse files Browse the repository at this point in the history
  • Loading branch information
xjules committed Sep 12, 2024
1 parent 3259be3 commit 8724f51
Showing 1 changed file with 48 additions and 0 deletions.
48 changes: 48 additions & 0 deletions tests/unit_tests/gui/plottery/test_ensemble_plot.py
Original file line number Diff line number Diff line change
@@ -0,0 +1,48 @@
from unittest.mock import Mock, patch

import pandas as pd
import pytest
from matplotlib.figure import Figure

from ert.gui.plottery import PlotConfig, PlotContext
from ert.gui.plottery.plots.ensemble import EnsemblePlot
from ert.gui.tools.plot.plot_api import EnsembleObject
from ert.shared.storage.summary_key_utils import is_rate


@pytest.fixture(
params=[
pytest.param("WOPR:OP_4"),
pytest.param("BPR:123"),
]
)
def plot_context(request):
context = Mock(spec=PlotContext)
context.ensembles.return_value = [
EnsembleObject("ensemble_1", "id", False, "experiment_1")
]
context.key.return_value = request.param
context.history_data = None
context.plotConfig.return_value = PlotConfig(title="Ensemble Plot")
return context


def test_ensemble_plot_handles_rate(plot_context: PlotContext):
figure = Figure()
with patch(
"ert.gui.plottery.plots.ensemble.EnsemblePlot._plotLines"
) as mock_plotLines:
EnsemblePlot().plot(
figure,
plot_context,
dict.fromkeys(
plot_context.ensembles(),
pd.DataFrame([[0.1], [0.2], [0.3], [0.4], [0.5]]),
),
pd.DataFrame(),
{},
)
if is_rate(plot_context.key()):
assert mock_plotLines.call_args[0][4] == "steps-pre"
else:
assert mock_plotLines.call_args[0][4] is None

0 comments on commit 8724f51

Please sign in to comment.