fix callback cancelling for partition, add test

Author: roveo

streamz/core.py

@@ -1025,14 +1025,15 @@ def update(self, x, who=None, metadata=None):
             metadata_buffer.extend(metadata)
         else:
             metadata_buffer.append(metadata)
+        if len(buffer) == self.n:
+            if self._timeout is not None and self.n > 1:
+                self._callbacks[key].cancel()
+            yield self._flush(key)
+            return
         if len(buffer) == 1 and self._timeout is not None:
             self._callbacks[key] = self.loop.call_later(
                 self._timeout, self._flush, key
             )
-        if len(buffer) == self.n:
-            if self._timeout is not None:
-                self._callbacks[key].cancel()
-            yield self._flush(key)
 
 
 @Stream.register_api()

streamz/tests/test_core.py

@@ -214,6 +214,15 @@ def test_partition_key_callable():
     assert L == [(0, 2), (1, 3), (4, 6), (5, 7)]
 
 
+def test_partition_size_one():
+    source = Stream()
+
+    source.partition(1, timeout=.01).sink(lambda x: None)
+
+    for i in range(10):
+        source.emit(i)
+
+
 def test_sliding_window():
     source = Stream()
     L = source.sliding_window(2).sink_to_list()