fix typing and tests

Author: erlenlh

src/ert/run_models/model_factory.py

@@ -322,25 +322,19 @@ def _setup_manual_update_enif(
     status_queue: SimpleQueue[StatusEvents],
 ) -> ManualUpdate:
     active_realizations = _realizations(args, config.runpath_config.num_realizations)
+    validate_minimum_realizations(config, active_realizations.tolist())
 
-    return ManualUpdateEnIF(
+    runmodel_config = ManualUpdateConfig(
         random_seed=config.random_seed,
         active_realizations=active_realizations.tolist(),
         ensemble_id=args.ensemble_id,
         minimum_required_realizations=config.analysis_config.minimum_required_realizations,
         target_ensemble=args.target_ensemble,
-        config=config,
         storage_path=config.ens_path,
         queue_config=config.queue_config,
         analysis_settings=config.analysis_config.es_settings,
         update_settings=update_settings,
-        status_queue=status_queue,
         runpath_file=config.runpath_file,
-        design_matrix=config.analysis_config.design_matrix,
-        parameter_configuration=config.ensemble_config.parameter_configuration,
-        response_configuration=config.ensemble_config.response_configuration,
-        derived_response_configuration=config.ensemble_config.derived_response_configuration,
-        ert_templates=config.ert_templates,
         user_config_file=Path(config.user_config_file),
         env_vars=config.env_vars,
         env_pr_fm_step=config.env_pr_fm_step,
@@ -349,9 +343,10 @@ def _setup_manual_update_enif(
         substitutions=config.substitutions,
         hooked_workflows=config.hooked_workflows,
         log_path=config.analysis_config.log_path,
-        observations=config.observation_declarations,
+        ert_templates=config.ert_templates,
         shape_registry=config.shape_registry,
     )
+    return ManualUpdateEnIF(**runmodel_config.model_dump(), status_queue=status_queue)
 
 
 def _setup_ensemble_smoother(

src/ert/run_models/run_model.py

@@ -23,6 +23,7 @@
 
 import numpy as np
 from pydantic import (
+    Field,
     PrivateAttr,
     computed_field,
     field_validator,
@@ -212,6 +213,8 @@ def _restore_plugin_forward_model_step_subclasses(self) -> Self:
 
 
 class RunModel(RunModelConfig, ABC):
+    status_queue: queue.SimpleQueue[StatusEvents] = Field(exclude=True, repr=False)
+
     # Private attributes initialized in model_post_init
     _initial_realizations_mask: list[bool] = PrivateAttr()
     _completed_realizations_mask: list[bool] = PrivateAttr(default_factory=list)
@@ -230,12 +233,10 @@ class RunModel(RunModelConfig, ABC):
     def __init__(
         self,
         *,
-        status_queue: queue.SimpleQueue[StatusEvents],
         _total_iterations: int | None = None,
         **data: Any,
     ) -> None:
         super().__init__(**data)
-        self._status_queue = status_queue
 
         if _total_iterations is not None:
             self._total_iterations = _total_iterations
@@ -276,7 +277,7 @@ def log_at_startup(self) -> None:
         keys_to_drop = [
             "_end_event",
             "_queue_config",
-            "_status_queue",
+            "status_queue",
             "_storage",
             "rng",
             "run_paths",
@@ -367,7 +368,7 @@ def experiment_type(self) -> ExperimentType:
         return self.__class__._experiment_type()
 
     def send_event(self, event: StatusEvents) -> None:
-        self._status_queue.put(event)
+        self.status_queue.put(event)
 
     @property
     def ensemble_size(self) -> int:

tests/ert/unit_tests/run_models/test_base_run_model.py

@@ -69,12 +69,12 @@ def test_run_model_does_not_support_rerun_failed_realizations(minimum_case):
 
 def test_status_when_rerunning_on_non_rerunnable_model(use_tmpdir):
     brm = create_run_model()
-    brm._status_queue = SimpleQueue()
+    brm.status_queue = SimpleQueue()
     brm.start_simulations_thread(
         EvaluatorServerConfig(use_token=False), rerun_failed_realizations=True
     )
-    assert isinstance(brm._status_queue.get(), StartEvent)
-    assert brm._status_queue.get() == EndEvent(
+    assert isinstance(brm.status_queue.get(), StartEvent)
+    assert brm.status_queue.get() == EndEvent(
         event_type="EndEvent",
         failed=True,
         msg="Run model None does not support restart/rerun of failed simulations.\n",

tests/ert/unit_tests/test_tracking.py

@@ -5,6 +5,7 @@
 from argparse import ArgumentParser
 from collections.abc import Generator
 from pathlib import Path
+from queue import SimpleQueue
 
 import pytest
 from jsonpath_ng import parse
@@ -34,8 +35,9 @@
 from ert.run_models.model_factory import create_model
 
 
-class Events:
+class Events(SimpleQueue):
     def __init__(self) -> None:
+        super().__init__()
         self.events: list[BaseModel] = []
         self.environment = []