1 files changed, 26 insertions, 20 deletions

diff --git a/mastodon/streaming.py b/mastodon/streaming.py
index d55ad54..1c73f48 100644
--- a/mastodon/streaming.py
+++ b/mastodon/streaming.py
@@ -6,7 +6,8 @@ https://github.com/tootsuite/mastodon/blob/master/docs/Using-the-API/Streaming-A
 import json
 import six
 from mastodon import Mastodon
-from mastodon.Mastodon import MastodonMalformedEventError
+from mastodon.Mastodon import MastodonMalformedEventError, MastodonNetworkError
+from requests.exceptions import ChunkedEncodingError
 
 class StreamListener(object):
     """Callbacks for the streaming API. Create a subclass, override the on_xxx
@@ -43,25 +44,31 @@ class StreamListener(object):
         """
         event = {}
         line_buffer = bytearray()
-        for chunk in response.iter_content(chunk_size = 1):
-            if chunk:
-                if chunk == b'\n':
-                    try:
-                        line = line_buffer.decode('utf-8')
-                    except UnicodeDecodeError as err:
-                        six.raise_from(
-                            MastodonMalformedEventError("Malformed UTF-8"),
-                            err
-                        )
-                    if line == '':
-                        self._dispatch(event)
-                        event = {}
+        try:
+            for chunk in response.iter_content(chunk_size = 1):
+                if chunk:
+                    if chunk == b'\n':
+                        try:
+                            line = line_buffer.decode('utf-8')
+                        except UnicodeDecodeError as err:
+                            six.raise_from(
+                                MastodonMalformedEventError("Malformed UTF-8"),
+                                err
+                            )
+                        if line == '':
+                            self._dispatch(event)
+                            event = {}
+                        else:
+                            event = self._parse_line(line, event)
+                        line_buffer = bytearray()
                     else:
-                        event = self._parse_line(line, event)
-                    line_buffer = bytearray()
-                else:
-                    line_buffer.extend(chunk)
-        
+                        line_buffer.extend(chunk)
+        except ChunkedEncodingError as err:
+            six.raise_from(
+                MastodonNetworkError("Server ceased communication."),
+                err
+            )
+            
     def _parse_line(self, line, event):
         if line.startswith(':'):
             self.handle_heartbeat()
@@ -102,7 +109,6 @@ class StreamListener(object):
                err
             )
         else:
-            # TODO: allow handlers to return/raise to stop streaming cleanly
             handler(payload)
 
 class CallbackStreamListener(StreamListener):