diff --git a/mitogen/core.py b/mitogen/core.py index 97c12bcd..d6a2fbb4 100644 --- a/mitogen/core.py +++ b/mitogen/core.py @@ -894,8 +894,9 @@ class Receiver(object): until :class:`ChannelError` is raised. """ while True: - msg = self.get(throw_dead=False) - if msg.is_dead: + try: + msg = self.get() + except ChannelError: return yield msg diff --git a/tests/receiver_test.py b/tests/receiver_test.py index 550b4525..8942ba29 100644 --- a/tests/receiver_test.py +++ b/tests/receiver_test.py @@ -37,6 +37,27 @@ class IterationTest(testlib.RouterMixin, testlib.TestCase): self.assertEquals(list(range(5)), list(m.unpickle() for m in recv)) self.assertEquals(10, ret.get().unpickle()) + def iter_and_put(self, recv, latch): + try: + for msg in recv: + latch.put(msg) + except Exception: + latch.put(sys.exc_info()[1]) + + def test_close_stops_iteration(self): + recv = mitogen.core.Receiver(self.router) + latch = mitogen.core.Latch() + t = threading.Thread( + target=self.iter_and_put, + args=(recv, latch), + ) + t.start() + t.join(0.1) + recv.close() + t.join() + self.assertTrue(latch.empty()) + + class CloseTest(testlib.RouterMixin, testlib.TestCase): def wait(self, latch, wait_recv):