Make zone_status_iterator raise a ServerNotReadyError when it gets an incomplete read

Author: BryceBeagle

brainframe/api/stubs/zone_statuses.py

@@ -1,11 +1,12 @@
+import json
 from typing import Dict, Generator
 
-import json
+import requests
 
-from brainframe.api.bf_codecs import ZoneStatus
+from brainframe.api import bf_codecs, bf_errors
 from .base_stub import BaseStub, DEFAULT_TIMEOUT
 
-ZONE_STATUS_TYPE = Dict[int, Dict[str, ZoneStatus]]
+ZONE_STATUS_TYPE = Dict[int, Dict[str, bf_codecs.ZoneStatus]]
 ZONE_STATUS_STREAM_TYPE = Generator[ZONE_STATUS_TYPE, None, None]
 
 
@@ -28,7 +29,7 @@ def get_latest_zone_statuses(self,
         data, _ = self._get_json(req, timeout)
 
         # Convert ZoneStatuses to Codecs
-        out = {int(s_id): {key: ZoneStatus.from_dict(val)
+        out = {int(s_id): {key: bf_codecs.ZoneStatus.from_dict(val)
                            for key, val in statuses.items()}
                for s_id, statuses in data.items()}
         return out
@@ -48,16 +49,26 @@ def get_zone_status_stream(self, timeout=None) -> ZONE_STATUS_STREAM_TYPE:
         def zone_status_iterator():
             # Don't use a timeout for this request, since it's ongoing
             resp = self._get(req, timeout=timeout)
-            for packet in resp.iter_lines(delimiter=b"\r\n"):
+
+            packets = resp.iter_lines(delimiter=b"\r\n")
+            while True:
+                try:
+                    packet = next(packets)
+                except requests.exceptions.ChunkedEncodingError as exc:
+                    message = "Incomplete packet while attempting to read " \
+                              "from zone status iterator"
+                    raise bf_errors.ServerNotReadyError(message) from exc
+
                 if packet == b'':
                     continue
 
                 # Parse the line
                 zone_statuses_dict = json.loads(packet)
 
-                processed = {int(s_id): {key: ZoneStatus.from_dict(val)
-                                         for key, val in statuses.items()}
-                             for s_id, statuses in zone_statuses_dict.items()}
+                processed = {
+                    int(s_id): {key: bf_codecs.ZoneStatus.from_dict(val)
+                                for key, val in statuses.items()}
+                    for s_id, statuses in zone_statuses_dict.items()}
                 yield processed
 
         return zone_status_iterator()