diff --git a/src/_ert/forward_model_runner/client.py b/src/_ert/forward_model_runner/client.py index ae28c919659..f98b7449af1 100644 --- a/src/_ert/forward_model_runner/client.py +++ b/src/_ert/forward_model_runner/client.py @@ -29,32 +29,6 @@ class Client: DEFAULT_ACK_TIMEOUT = 5 _receiver_task: asyncio.Task[None] | None - async def __aenter__(self) -> Self: - await self.connect() - return self - - async def __aexit__( - self, exc_type: Any, exc_value: Any, exc_traceback: Any - ) -> None: - try: - await self.send(DISCONNECT_MSG) - except ClientConnectionError: - logger.error("No ack for dealer disconnection. Connection is down!") - finally: - self.socket.disconnect(self.url) - await self._term_receiver_task() - self.term() - - def term(self) -> None: - self.socket.close() - self.context.term() - - async def _term_receiver_task(self) -> None: - if self._receiver_task and not self._receiver_task.done(): - self._receiver_task.cancel() - await asyncio.gather(self._receiver_task, return_exceptions=True) - self._receiver_task = None - def __init__( self, url: str, @@ -83,6 +57,32 @@ def __init__( self._receiver_task = None + async def __aenter__(self) -> Self: + await self.connect() + return self + + async def __aexit__( + self, exc_type: Any, exc_value: Any, exc_traceback: Any + ) -> None: + try: + await self.send(DISCONNECT_MSG) + except ClientConnectionError: + logger.error("No ack for dealer disconnection. Connection is down!") + finally: + self.socket.disconnect(self.url) + await self._term_receiver_task() + self.term() + + def term(self) -> None: + self.socket.close() + self.context.term() + + async def _term_receiver_task(self) -> None: + if self._receiver_task and not self._receiver_task.done(): + self._receiver_task.cancel() + await asyncio.gather(self._receiver_task, return_exceptions=True) + self._receiver_task = None + async def connect(self) -> None: self.socket.connect(self.url) await self._term_receiver_task()