Refactor to common mockstorage in gui tests

Author: jonathan-eq

src/ert/gui/simulation/manual_update_panel.py

@@ -120,8 +120,10 @@ def _realizations_from_fs(self) -> None:
                 parameters = ensemble.get_realization_mask_with_parameters()
                 responses = ensemble.get_realization_mask_with_responses()
                 mask = np.logical_and(parameters, responses)
-                self._ensemble_size = list(mask).count(True)
-
+                self._number_of_realizations_with_responses_in_selected_ensemble = list(
+                    mask
+                ).count(True)
+                self._ensemble_size = ensemble.ensemble_size
                 self._active_realizations_field.setValidator(
                     EnsembleRealizationsArgument(
                         lambda: ensemble,
@@ -135,7 +137,7 @@ def _realizations_from_fs(self) -> None:
                 self._active_realizations_model.ensemble_size = ensemble.ensemble_size
                 self._active_realizations_model.setValueFromMask(mask)
                 self._number_of_realizations_label.setText(
-                    f"<b>{self._ensemble_size}</b>"
+                    f"<b>{self._number_of_realizations_with_responses_in_selected_ensemble}</b>"
                 )
                 self._analysis_module_edit.ensemble_size = ensemble.ensemble_size
                 self._analysis_module_edit.setEnabled(bool(ensemble.ensemble_size))

tests/ert/ui_tests/gui/test_restart_esmda.py

@@ -66,7 +66,7 @@ def test_active_realizations_esmda(opened_main_window_poly, qtbot):
     simulation_mode_combo.setCurrentText(MultipleDataAssimilation.display_name())
     es_mda_panel = gui.findChild(QWidget, name="ES_MDA_panel")
     assert es_mda_panel
-    active_reals = gui.findChild(StringBox, "active_realizations_box")
+    active_reals = es_mda_panel.findChild(StringBox, "active_realizations_box")
     assert active_reals.text() == "0-9"
 
     restart_checkbox = es_mda_panel.findChild(QCheckBox, name="restart_checkbox_esmda")

tests/ert/unit_tests/gui/simulation/conftest.py

@@ -0,0 +1,66 @@
+from datetime import datetime
+from uuid import uuid4
+
+from ert.storage.local_ensemble import LocalEnsemble
+from ert.storage.local_ensemble import _Index as _EnsembleIndex
+from ert.storage.local_experiment import LocalExperiment
+from ert.storage.local_experiment import _Index as _ExperimentIndex
+from ert.storage.local_storage import LocalStorage
+from ert.storage.realization_storage_state import RealizationStorageState
+
+REALIZATION_FINISHED_SUCCESSFULLY = {
+    RealizationStorageState.PARAMETERS_LOADED,
+    RealizationStorageState.RESPONSES_LOADED,
+}
+REALIZATION_UNDEFINED = {RealizationStorageState.UNDEFINED}
+REALIZATION_ONLY_PARAMETERS = {RealizationStorageState.PARAMETERS_LOADED}
+REALIZATION_FAILED_DURING_EVALUATION = {
+    RealizationStorageState.PARAMETERS_LOADED,
+    RealizationStorageState.FAILURE_IN_CURRENT,
+}
+
+
+class MockEnsemble(LocalEnsemble):
+    def __init__(self, ensemble_name, experiment_id, storage_states, storage) -> None:
+        self._index = _EnsembleIndex(
+            id=uuid4(),
+            experiment_id=experiment_id,
+            ensemble_size=len(storage_states),
+            iteration=0,
+            name=ensemble_name,
+            prior_ensemble_id=None,
+            started_at=datetime.now(),
+        )
+        self._storage_state = storage_states
+        self._storage = storage
+
+    def get_ensemble_state(self):
+        return self._storage_state
+
+
+class MockExperiment(LocalExperiment):
+    def __init__(self, experiment_name) -> None:
+        self._index = _ExperimentIndex(id=uuid4(), name=experiment_name)
+
+    @property
+    def relative_weights(self) -> str:
+        return "4, 2, 1"
+
+
+class MockStorage(LocalStorage):
+    def __init__(self) -> None:
+        self._ensembles = {}
+        self._experiments = {}
+
+    def _setup_mocked_run(
+        self, ensemble_name, experiment_name, ensemble_states
+    ) -> None:
+        mock_experiment = MockExperiment(experiment_name)
+        mock_ensemble2 = MockEnsemble(
+            ensemble_name,
+            experiment_id=mock_experiment.id,
+            storage_states=ensemble_states,
+            storage=self,
+        )
+        self._ensembles[mock_ensemble2.id] = mock_ensemble2
+        self._experiments[mock_experiment.id] = mock_experiment

tests/ert/unit_tests/gui/simulation/test_manual_update.py

@@ -1,9 +1,3 @@
-import datetime
-from unittest.mock import MagicMock
-from uuid import uuid4
-
-import numpy as np
-import pytest
 from PyQt6.QtCore import Qt
 from pytestqt.qtbot import QtBot
 
@@ -12,87 +6,14 @@
 from ert.gui.ertwidgets.ensembleselector import EnsembleSelector
 from ert.gui.ertwidgets.stringbox import StringBox
 from ert.gui.simulation.manual_update_panel import ManualUpdatePanel
-from ert.storage.local_ensemble import LocalEnsemble
-from ert.storage.local_experiment import LocalExperiment
-from ert.storage.local_storage import LocalStorage
-from ert.storage.realization_storage_state import RealizationStorageState
-
-realization_finished_successfully = {
-    RealizationStorageState.PARAMETERS_LOADED,
-    RealizationStorageState.RESPONSES_LOADED,
-}
-realization_failed = {RealizationStorageState.UNDEFINED}
-
-
-class MockStorage(LocalStorage):
-    def __init__(self) -> None:
-        self._ensembles = {}
-        self._experiments = {}
-
-    def _setup_mocked_run(
-        self, ensemble_name, experiment_name, ensemble_states
-    ) -> None:
-        mock_ensemble = MagicMock(spec=LocalEnsemble)
-        mock_ensemble.id = uuid4()
-        mock_ensemble.name = ensemble_name
-        mock_ensemble.relative_weights = "4, 2, 1"
-        mock_ensemble.started_at = datetime.datetime.now()
-
-        def mock_get_ensemble_state(*args, **kwargs):
-            return ensemble_states
-
-        mock_ensemble.ensemble_size = len(mock_get_ensemble_state())
-        mock_ensemble.get_ensemble_state = mock_get_ensemble_state
-
-        def mock_get_realization_mask_with_responses(*args):
-            return np.array(
-                [
-                    RealizationStorageState.RESPONSES_LOADED in state
-                    for state in mock_ensemble.get_ensemble_state()
-                ]
-            )
-
-        def mock_get_realization_mask_with_parameters(*args):
-            return np.array(
-                [
-                    RealizationStorageState.PARAMETERS_LOADED in state
-                    for state in mock_ensemble.get_ensemble_state()
-                ]
-            )
-
-        mock_ensemble.get_realization_mask_with_responses = (
-            mock_get_realization_mask_with_responses
-        )
-        mock_ensemble.get_realization_mask_with_parameters = (
-            mock_get_realization_mask_with_parameters
-        )
 
-        def mock_get_realization_list_with_responses(self) -> list[int]:
-            mask = self.get_realization_mask_with_responses()
-            return np.where(mask)[0].tolist()
+from .conftest import (
+    REALIZATION_FINISHED_SUCCESSFULLY,
+    REALIZATION_UNDEFINED,
+    MockStorage,
+)
 
-        mock_ensemble.has_data = (
-            lambda *args: len(mock_ensemble.get_realization_list_with_responses()) > 0
-        )
-        mock_ensemble.get_realization_list_with_responses = (
-            lambda *args: mock_get_realization_list_with_responses(mock_ensemble)
-        )
-        mock_experiment = MagicMock(spec=LocalExperiment)
-        mock_experiment.ensembles = [mock_ensemble]
-        mock_experiment.id = uuid4()
-        mock_experiment.name = experiment_name
 
-        mock_ensemble.experiment_id = mock_experiment.id
-        mock_ensemble.experiment = mock_experiment
-        self._ensembles[mock_ensemble.id] = mock_ensemble
-        self._experiments[mock_experiment.id] = mock_experiment
-
-    @property
-    def ensembles(self):
-        yield from self._ensembles.values()
-
-
-@pytest.mark.usefixtures("copy_poly_case")
 def test_that_active_realizations_selector_validates_with_ensemble_size_from_prior(
     qtbot: QtBot,
 ) -> None:
@@ -107,26 +28,26 @@ def test_that_active_realizations_selector_validates_with_ensemble_size_from_pri
         "mock_ensemble0",
         "mock_experiment0",
         [
-            realization_failed,
-            realization_finished_successfully,
-            realization_finished_successfully,
-            realization_finished_successfully,
-            realization_failed,
-            realization_finished_successfully,
-            realization_finished_successfully,
-            realization_finished_successfully,
+            REALIZATION_UNDEFINED,
+            REALIZATION_FINISHED_SUCCESSFULLY,
+            REALIZATION_FINISHED_SUCCESSFULLY,
+            REALIZATION_FINISHED_SUCCESSFULLY,
+            REALIZATION_UNDEFINED,
+            REALIZATION_FINISHED_SUCCESSFULLY,
+            REALIZATION_FINISHED_SUCCESSFULLY,
+            REALIZATION_FINISHED_SUCCESSFULLY,
         ],
     )
     notifier._storage._setup_mocked_run(
         "mock_ensemble1",
         "mock_experiment1",
         [
-            realization_failed,
-            realization_failed,
-            realization_finished_successfully,
-            realization_finished_successfully,
-            realization_failed,
-            realization_finished_successfully,
+            REALIZATION_UNDEFINED,
+            REALIZATION_UNDEFINED,
+            REALIZATION_FINISHED_SUCCESSFULLY,
+            REALIZATION_FINISHED_SUCCESSFULLY,
+            REALIZATION_UNDEFINED,
+            REALIZATION_FINISHED_SUCCESSFULLY,
         ],
     )
     panel = ManualUpdatePanel(
@@ -138,8 +59,10 @@ def test_that_active_realizations_selector_validates_with_ensemble_size_from_pri
     qtbot.addWidget(panel)
 
     prior_ensemble0_number_of_realizations_with_responses = 6
-    prior_ensemble1_number_of_realizations_with_responses = 3
+    prior_ensemble0_ensemble_size = 8
 
+    prior_ensemble1_number_of_realizations_with_responses = 3
+    prior_ensemble1_ensemble_size = 6
     realization_selector = panel.findChild(StringBox, "active_realizations_box")
     assert realization_selector is not None
     ensemble_selector = panel.findChild(EnsembleSelector)
@@ -159,6 +82,9 @@ def test_that_active_realizations_selector_validates_with_ensemble_size_from_pri
         realization_selector._validation._validation_message
     )
     assert panel.isConfigurationValid()
+    assert (
+        panel.get_experiment_arguments().ensemble_size == prior_ensemble0_ensemble_size
+    )
 
     # We try running a realization that does not have RESPONSES_LOADED
     realization_selector.setText("4-7")
@@ -169,14 +95,16 @@ def test_that_active_realizations_selector_validates_with_ensemble_size_from_pri
 
     # The active realizations field should auto-populate with a valid value
     assert panel.isConfigurationValid()
+    assert (
+        panel.get_experiment_arguments().ensemble_size == prior_ensemble1_ensemble_size
+    )
     assert (
         panel._number_of_realizations_label.text()
         == f"<b>{prior_ensemble1_number_of_realizations_with_responses}</b>"
     )
     assert realization_selector.text() == "2-3, 5"
 
 
-@pytest.mark.usefixtures("copy_poly_case")
 def test_that_manual_update_ensemble_selector_only_shows_ensembles_with_data(
     qtbot: QtBot,
 ) -> None:
@@ -190,7 +118,7 @@ def test_that_manual_update_ensemble_selector_only_shows_ensembles_with_data(
     notifier._storage._setup_mocked_run(
         "mock_ensemble_no_data",
         "mock_experiment2",
-        [realization_failed, realization_failed, realization_failed],
+        [REALIZATION_UNDEFINED, REALIZATION_UNDEFINED, REALIZATION_UNDEFINED],
     )
 
     panel = ManualUpdatePanel(