Fix the callback shutting itself off

The background task can't return obviously if we want the stream to
continue operating.

Author: dwskoog

streamz/core.py

@@ -752,25 +752,26 @@ def __init__(self, upstream, func, *args, **kwargs):
         stream_name = kwargs.pop('stream_name', None)
         self.kwargs = kwargs
         self.args = args
-        self.work_queue = asyncio.Queue()
+        self.input_queue = asyncio.Queue()
 
         Stream.__init__(self, upstream, stream_name=stream_name, ensure_io_loop=True)
-        self.cb_task = self.loop.asyncio_loop.create_task(self.cb())
+        self.input_task = self.loop.asyncio_loop.create_task(self.input_callback())
 
     def update(self, x, who=None, metadata=None):
         coro = self.func(x, *self.args, **self.kwargs)
-        return self.work_queue.put_nowait((coro, metadata))
+        self.input_queue.put_nowait((coro, metadata))
 
-    async def cb(self):
+    async def input_callback(self):
         while True:
-            coro, metadata = await self.work_queue.get()
             try:
+                coro, metadata = await self.input_queue.get()
+                self.input_queue.task_done()
                 result = await coro
             except Exception as e:
                 logger.exception(e)
                 raise
             else:
-                return self._emit(result, metadata=metadata)
+                self._emit(result, metadata=metadata)
 
 
 @Stream.register_api()

streamz/tests/test_core.py

@@ -139,8 +139,13 @@ async def add_native(x=0, y=0):
     L = source.map_async(add_tor, y=1).map_async(add_native, y=2).sink_to_list()
 
     yield source.emit(0)
+    yield source.emit(1)
+    yield source.emit(2)
 
-    yield await_for(lambda: L == [3], 1)
+    def fail_func():
+        assert L == [3, 4, 5]
+
+    yield await_for(lambda: L == [3, 4, 5], 1, fail_func=fail_func)
 
 
 @pytest.mark.asyncio
@@ -156,8 +161,13 @@ async def add_native(x=0, y=0):
     L = source.map_async(add_tor, y=1).map_async(add_native, y=2).sink_to_list()
 
     await source.emit(0)
+    await source.emit(1)
+    await source.emit(2)
+
+    def fail_func():
+        assert L == [3, 4, 5]
 
-    await await_for(lambda: L == [3], 1)
+    await await_for(lambda: L == [3, 4, 5], 1, fail_func=fail_func)
 
 
 def test_map_args():