From 79588cd0220c560cd786f99a31b08c269501337c Mon Sep 17 00:00:00 2001
From: Julius Parulek <parulek@gmail.com>
Date: Tue, 12 Dec 2023 16:27:49 +0100
Subject: [PATCH] --amend

---
 src/ert/scheduler/job.py                     | 3 +++
 src/ert/scheduler/local_driver.py            | 4 ++++
 tests/unit_tests/scheduler/test_scheduler.py | 3 +--
 3 files changed, 8 insertions(+), 2 deletions(-)

diff --git a/src/ert/scheduler/job.py b/src/ert/scheduler/job.py
index 19534f261dc..9798ebf6a50 100644
--- a/src/ert/scheduler/job.py
+++ b/src/ert/scheduler/job.py
@@ -85,6 +85,9 @@ async def __call__(
 
                 await self._send(State.RUNNING)
                 returncode = await self.returncode
+                # we need to make sure that the task has finished too
+                await self.driver.wait_to_finish(self.real.iens)
+
                 if (
                     returncode == 0
                     and forward_model_ok(self.real.run_arg).status
diff --git a/src/ert/scheduler/local_driver.py b/src/ert/scheduler/local_driver.py
index 960d98e9bc4..542798f1011 100644
--- a/src/ert/scheduler/local_driver.py
+++ b/src/ert/scheduler/local_driver.py
@@ -23,6 +23,10 @@ async def kill(self, iens: int) -> None:
         except KeyError:
             return
 
+    async def wait_to_finish(self, iens: int):
+        # we might need to do some timeout here I guess
+        await self._tasks[iens]
+
     async def _wait_until_finish(
         self, iens: int, executable: str, /, *args: str, cwd: str
     ) -> None:
diff --git a/tests/unit_tests/scheduler/test_scheduler.py b/tests/unit_tests/scheduler/test_scheduler.py
index 9224b6d70da..59f7159acbc 100644
--- a/tests/unit_tests/scheduler/test_scheduler.py
+++ b/tests/unit_tests/scheduler/test_scheduler.py
@@ -121,6 +121,5 @@ async def test_that_max_submit_was_reached(tmp_path: Path, realization):
     sch.add_realization(realization, callback_timeout=lambda _: None)
     create_jobs_json(tmp_path, [step])
     sch.add_dispatch_information_to_jobs_file()
-    scheduler_task = asyncio.create_task(sch.execute())
-    await scheduler_task
+    assert await sch.execute() == EVTYPE_ENSEMBLE_STOPPED
     assert (tmp_path / "cnt").read_text() == "3\n"