Fix regression from async forward_model

equinor · Nov 24, 2023 · 93c134d · 93c134d
1 parent 82ccc29
commit 93c134d
Show file tree

Hide file tree

Showing 3 changed files with 8 additions and 6 deletions.
diff --git a/tests/unit_tests/ensemble_evaluator/conftest.py b/tests/unit_tests/ensemble_evaluator/conftest.py
@@ -2,7 +2,7 @@
 import os
 import stat
 from pathlib import Path
-from unittest.mock import MagicMock, Mock
+from unittest.mock import AsyncMock, MagicMock, Mock
 
 import pytest
 
@@ -64,9 +64,10 @@ def queue_config_fixture():
 @pytest.fixture
 def make_ensemble_builder(queue_config):
     def _make_ensemble_builder(monkeypatch, tmpdir, num_reals, num_jobs, job_sleep=0):
+        forward_mock = AsyncMock(return_value=(LoadStatus.LOAD_SUCCESSFUL, ""))
         monkeypatch.setattr(
             "ert.job_queue.queue.forward_model_ok",
-            lambda _: (LoadStatus.LOAD_SUCCESSFUL, ""),
+            AsyncMock(return_value=(LoadStatus.LOAD_SUCCESSFUL, "")),
         )
         builder = ert.ensemble_evaluator.EnsembleBuilder()
         with tmpdir.as_cwd():

diff --git a/tests/unit_tests/job_queue/conftest.py b/tests/unit_tests/job_queue/conftest.py
@@ -1,6 +1,6 @@
 import stat
 from pathlib import Path
-from unittest.mock import MagicMock
+from unittest.mock import AsyncMock
 
 import pytest
 
@@ -10,7 +10,7 @@
 
 @pytest.fixture
 def mock_fm_ok(monkeypatch):
-    fm_ok = MagicMock(return_value=(LoadStatus.LOAD_SUCCESSFUL, ""))
+    fm_ok = AsyncMock(return_value=(LoadStatus.LOAD_SUCCESSFUL, ""))
     monkeypatch.setattr(ert.job_queue.queue, "forward_model_ok", fm_ok)
     yield fm_ok
 

diff --git a/tests/unit_tests/job_queue/test_job_queue.py b/tests/unit_tests/job_queue/test_job_queue.py
@@ -6,7 +6,7 @@
 from pathlib import Path
 from threading import BoundedSemaphore
 from typing import Any, Callable, Dict, List, Optional
-from unittest.mock import MagicMock, patch
+from unittest.mock import AsyncMock, MagicMock, patch
 
 import pytest
 
@@ -195,7 +195,8 @@ async def test_run_done_callback(
 ):
     monkeypatch.chdir(tmpdir)
     monkeypatch.setattr(
-        "ert.job_queue.queue.forward_model_ok", lambda _: (loadstatus, "foo")
+        "ert.job_queue.queue.forward_model_ok",
+        AsyncMock(return_value=(loadstatus, "foo")),
     )
     job_queue = create_local_queue(simple_script)
     execute_task = asyncio.create_task(job_queue.execute())