From 822978520f625958e828e1c6ade23cc2e582b570 Mon Sep 17 00:00:00 2001
From: David Wilson <dw@botanicus.net>
Date: Sun, 20 Jan 2019 08:13:04 +0000
Subject: [PATCH] issue #446: update Receiver.__iter__ to match

iter() previously relied on the fake dead message being enqueued.
---
 mitogen/core.py        |  5 +++--
 tests/receiver_test.py | 21 +++++++++++++++++++++
 2 files changed, 24 insertions(+), 2 deletions(-)

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):