diff --git a/vrtManager/connection.py b/vrtManager/connection.py index 6eb7349e..8689cfbb 100644 --- a/vrtManager/connection.py +++ b/vrtManager/connection.py @@ -147,9 +147,11 @@ def other_hosts(): then goes on with the host it used before, through the objects it kept: the request counts on that host again after the block (HOST_REQUESTS).""" held = getattr(_request, "held", None) - yield - if held is not None: - _count_on(held) + try: + yield + finally: + if held is not None: + _count_on(held) request_started.connect(_request_started, dispatch_uid="vrtManager.connection.request_started") diff --git a/vrtManager/tests/test_connection.py b/vrtManager/tests/test_connection.py index fb68e40e..f1eed0c8 100644 --- a/vrtManager/tests/test_connection.py +++ b/vrtManager/tests/test_connection.py @@ -1,5 +1,6 @@ """Reachability probe and connections that do not block each other.""" +import contextlib import os import socket import threading @@ -296,11 +297,19 @@ class HostRequestsTestCase(unittest.TestCase): self.assertEqual(self.request("h") + self.request("h"), []) def test_a_request_back_from_other_hosts_counts_on_its_own_again(self): + self.back_from_other_hosts(fails=False) + + def test_a_request_back_from_other_hosts_that_failed_counts_on_its_own_again(self): + self.back_from_other_hosts(fails=True) + + def back_from_other_hosts(self, fails): def run(): request_started.send(sender=None) self.manager.get_connection("h", "u", "p", CONN_TCP) - with connection.other_hosts(): + with contextlib.suppress(RuntimeError), connection.other_hosts(): self.manager.get_connection("g", "u", "p", CONN_TCP) + if fails: + raise RuntimeError("the other host's read failed") back.set() self.done.wait(5) request_finished.send(sender=None)