Fixing existing tests and linting complaints

Author: freol35241

streamz/core.py

@@ -255,7 +255,7 @@ def __init__(self, upstream=None, upstreams=None, stream_name=None,
             self.upstreams = []
 
         # Lazily loaded exception handler to avoid recursion
-        self._on_exception =  None
+        self._on_exception = None
 
         self._set_asynchronous(asynchronous)
         self._set_loop(loop)
@@ -678,30 +678,36 @@ def _release_refs(self, metadata, n=1):
                 m['ref'].release(n)
 
     def on_exception(self):
-        """Returns the exception handler associated with this stream
+        """ Returns the exception handler associated with this stream. The exception handler is either lazily loaded
+        at this point or (if alredy loaded) just returned.
         """
         self._on_exception = self._on_exception or _on_exception()
         return self._on_exception
 
 
 class InvalidDataError(Exception):
-    pass
+    """Generic error that is raised when data passed into a node causes an exception
+    """
+
 
 class _on_exception(Stream):
+    """ Internal exception-handler for Stream-nodes.
+    """
 
     def __init__(self, *args, **kwargs):
         self.silent = False
         Stream.__init__(self, *args, **kwargs)
 
     def update(self, x, who=None, metadata=None):
         cause, exc = x
-        
+
         if self.silent or len(self.downstreams) > 0:
             return self._emit(x, metadata=metadata)
         else:
             logger.exception(exc)
             raise InvalidDataError(cause) from exc
 
+
 @Stream.register_api()
 class map(Stream):
     """ Apply a function to every element in the stream
@@ -737,7 +743,7 @@ def __init__(self, upstream, func, *args, **kwargs):
 
     def update(self, x, who=None, metadata=None):
         result = self.func(x, *self.args, **self.kwargs)
-        self._emit(result, metadata=metadata)
+        return self._emit(result, metadata=metadata)
 
 
 @Stream.register_api()

streamz/tests/test_core.py

@@ -16,6 +16,7 @@
 import streamz as sz
 
 from streamz import RefCounter
+from streamz.core import InvalidDataError
 from streamz.sources import sink_to_file
 from streamz.utils_test import (inc, double, gen_test, tmpfile, captured_logger,   # noqa: F401
         clean, await_for, metadata, wait_for)  # noqa: F401
@@ -933,7 +934,7 @@ def test_pluck():
     assert L == [2]
     a.emit([4, 5, 6, 7, 8, 9])
     assert L == [2, 5]
-    with pytest.raises(IndexError):
+    with pytest.raises(InvalidDataError):
         a.emit([1])
 
 
@@ -945,7 +946,7 @@ def test_pluck_list():
     assert L == [(1, 3)]
     a.emit([4, 5, 6, 7, 8, 9])
     assert L == [(1, 3), (4, 6)]
-    with pytest.raises(IndexError):
+    with pytest.raises(InvalidDataError):
         a.emit([1])
 
 
@@ -1579,7 +1580,7 @@ def test_map_errors_log():
 def test_map_errors_raises():
     a = Stream()
     b = a.map(lambda x: 1 / x)  # noqa: F841
-    with pytest.raises(ZeroDivisionError):
+    with pytest.raises(InvalidDataError):
         a.emit(0)
 
 
@@ -1599,7 +1600,7 @@ def test_accumulate_errors_log():
 def test_accumulate_errors_raises():
     a = Stream()
     b = a.accumulate(lambda x, y: x / y, with_state=True)  # noqa: F841
-    with pytest.raises(ZeroDivisionError):
+    with pytest.raises(InvalidDataError):
         a.emit(1)
         a.emit(0)