From 3e33ecfefd91f9f4335d610fffc7d7c3900537b8 Mon Sep 17 00:00:00 2001 From: Martin Raspaud Date: Mon, 4 Dec 2023 16:01:31 +0100 Subject: [PATCH] Make patched_subscriber_recv interruptible This makes it possible to catch cases when the subscriber is closed before all expected messages have arrived. --- posttroll/testing.py | 11 ++++++++++- 1 file changed, 10 insertions(+), 1 deletion(-) diff --git a/posttroll/testing.py b/posttroll/testing.py index 501aa53..526d161 100644 --- a/posttroll/testing.py +++ b/posttroll/testing.py @@ -1,11 +1,20 @@ """Testing utilities.""" from contextlib import contextmanager + @contextmanager def patched_subscriber_recv(messages): """Patch the Subscriber object to return given messages.""" from unittest import mock - with mock.patch("posttroll.subscriber.Subscriber.recv", mock.Mock(return_value=messages)): + + def interuptible_recv(self): + """Yield message until the subscriber is closed.""" + for msg in messages: + if self._loop is False: + break + yield msg + + with mock.patch("posttroll.subscriber.Subscriber.recv", interuptible_recv): yield @contextmanager