From 353edda58fe965573ecc59f84cd51c7511edfe4f Mon Sep 17 00:00:00 2001
From: Jon Holba <jholba@equinor.com>
Date: Mon, 3 Jun 2024 13:07:01 +0200
Subject: [PATCH] Close Evaluator server for new connections earlier

Ensemble evaluator accepted new connections while shutting down before
this commit
---
 src/ert/ensemble_evaluator/evaluator.py       |  3 ++-
 .../test_ensemble_evaluator.py                | 22 +++++++++++++++++++
 2 files changed, 24 insertions(+), 1 deletion(-)

diff --git a/src/ert/ensemble_evaluator/evaluator.py b/src/ert/ensemble_evaluator/evaluator.py
index 0d21c54a72b..815be171dec 100644
--- a/src/ert/ensemble_evaluator/evaluator.py
+++ b/src/ert/ensemble_evaluator/evaluator.py
@@ -313,9 +313,10 @@ async def _server(self) -> None:
             ping_timeout=60,
             ping_interval=60,
             close_timeout=60,
-        ):
+        ) as server:
             self._server_started.set()
             await self._server_done.wait()
+            server.close(close_connections=False)
             if self._dispatchers_connected is not None:
                 logger.debug(
                     f"Got done signal. {self._dispatchers_connected.qsize()} "
diff --git a/tests/ert/unit_tests/ensemble_evaluator/test_ensemble_evaluator.py b/tests/ert/unit_tests/ensemble_evaluator/test_ensemble_evaluator.py
index 90e47f10d15..f4cf5cca058 100644
--- a/tests/ert/unit_tests/ensemble_evaluator/test_ensemble_evaluator.py
+++ b/tests/ert/unit_tests/ensemble_evaluator/test_ensemble_evaluator.py
@@ -94,6 +94,28 @@ async def mock_done_prematurely(message, *args, **kwargs):
         await evaluator.run_and_get_successful_realizations()
 
 
+async def test_new_connections_are_denied_when_evaluator_is_closing_down(
+    evaluator_to_use,
+):
+    evaluator = evaluator_to_use
+
+    class TestMonitor(Monitor):
+        def __init__(self, *args, **kwargs):
+            super().__init__(*args, **kwargs)
+            self._connection_timeout = 1
+
+    async def new_connection():
+        await evaluator._server_done.wait()
+        async with TestMonitor(evaluator._config.get_connection_info()):
+            pass
+
+    new_connection_task = asyncio.create_task(new_connection())
+    evaluator.stop()
+
+    with pytest.raises(RuntimeError):
+        await new_connection_task
+
+
 @pytest.fixture(name="evaluator_to_use")
 async def evaluator_to_use_fixture(make_ee_config):
     ensemble = TestEnsemble(0, 2, 2, id_="0")