From 511f613c3328d0fd3a5b6e38d3ab76be434bd9d2 Mon Sep 17 00:00:00 2001
From: xjules <jparu@equinor.com>
Date: Fri, 13 Sep 2024 09:12:51 +0200
Subject: [PATCH] Add EnsemblePlot support for rate plotting drawing

 - exploits "steps-pre" drawing style
 - tests that draw_style is set correctly
---
 src/ert/gui/plottery/plots/ensemble.py        |  7 ++-
 .../gui/plottery/test_ensemble_plot.py        | 48 +++++++++++++++++++
 2 files changed, 54 insertions(+), 1 deletion(-)
 create mode 100644 tests/unit_tests/gui/plottery/test_ensemble_plot.py

diff --git a/src/ert/gui/plottery/plots/ensemble.py b/src/ert/gui/plottery/plots/ensemble.py
index 57b308d663f..58dd3ab0d9a 100644
--- a/src/ert/gui/plottery/plots/ensemble.py
+++ b/src/ert/gui/plottery/plots/ensemble.py
@@ -1,12 +1,13 @@
 from __future__ import annotations
 
-from typing import TYPE_CHECKING, Dict
+from typing import TYPE_CHECKING, Dict, Optional
 
 import numpy as np
 import pandas as pd
 
 from ert.gui.plottery.plots.history import plotHistory
 from ert.gui.tools.plot.plot_api import EnsembleObject
+from ert.shared.storage.summary_key_utils import is_rate
 
 from .observations import plotObservations
 from .plot_tools import PlotTools
@@ -36,6 +37,7 @@ def plot(
 
         plot_context.y_axis = plot_context.VALUE_AXIS
         plot_context.x_axis = plot_context.DATE_AXIS
+        draw_style = "steps-pre" if is_rate(plot_context.key()) else None
 
         for ensemble, data in ensemble_to_data_map.items():
             data = data.T
@@ -50,6 +52,7 @@ def plot(
                     config,
                     data,
                     f"{ensemble.experiment_name} : {ensemble.name}",
+                    draw_style,
                 )
                 config.nextColor()
 
@@ -71,6 +74,7 @@ def _plotLines(
         plot_config: PlotConfig,
         data: pd.DataFrame,
         ensemble_label: str,
+        draw_style: Optional[str] = None,
     ) -> None:
         style = plot_config.defaultStyle()
 
@@ -86,6 +90,7 @@ def _plotLines(
             linewidth=style.width,
             linestyle=style.line_style,
             markersize=style.size,
+            drawstyle=draw_style,
         )
 
         if len(lines) > 0:
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