From 8724f51dfbfea6ebb35b76a8f22f5a5f6a5057d1 Mon Sep 17 00:00:00 2001
From: xjules <jparu@equinor.com>
Date: Thu, 12 Sep 2024 16:07:05 +0200
Subject: [PATCH] Add ensemble plot test for correct rate draw style

---
 .../gui/plottery/test_ensemble_plot.py        | 48 +++++++++++++++++++
 1 file changed, 48 insertions(+)
 create mode 100644 tests/unit_tests/gui/plottery/test_ensemble_plot.py

diff --git a/tests/unit_tests/gui/plottery/test_ensemble_plot.py b/tests/unit_tests/gui/plottery/test_ensemble_plot.py
new file mode 100644
index 00000000000..f6778a56b23
--- /dev/null
+++ b/tests/unit_tests/gui/plottery/test_ensemble_plot.py
@@ -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