From bbca6f603b5a7806c5d361a5f29a3c11e20aa2e4 Mon Sep 17 00:00:00 2001
From: Arthur Bols <arthur@bols.dev>
Date: Sat, 27 Mar 2021 19:08:06 +0100
Subject: [PATCH] Improve ChunkedRetriever error handling and documentation

---
 httplib/retriever.py | 31 ++++++++++++++++++-------------
 1 file changed, 18 insertions(+), 13 deletions(-)

diff --git a/httplib/retriever.py b/httplib/retriever.py
index f0c009f..3c7e35d 100644
--- a/httplib/retriever.py
+++ b/httplib/retriever.py
@@ -126,11 +126,9 @@ class ContentLengthRetriever(Retriever):
             try:
                 buffer = self.client.read(remaining)
             except TimeoutError:
-                logging.error("Timed out before receiving the complete payload")
                 raise IncompleteResponse("Timed out before receiving complete payload")
             except ConnectionError:
-                logging.error("Connection closed before receiving the complete payload")
-                raise IncompleteResponse("Connection closed before receiving complete payload")
+                raise IncompleteResponse("Connection closed before receiving the complete payload")
 
             if len(buffer) == 0:
                 logging.warning("Received payload length %s less than expected %s", cur_payload_size, self.length)
@@ -164,19 +162,26 @@ class ChunkedRetriever(Retriever):
         Returns an iterator of the received message bytes.
         The size of each iteration is not necessarily constant.
         @raise IncompleteResponse: if the connection is closed or timed out before receiving the complete payload.
+        @raise InvalidResponse: if the length of a chunk could not be determined.
         """
-        while True:
-            chunk_size = self.__get_chunk_size()
-            logging.debug("chunk-size: %s", chunk_size)
-            if chunk_size == 0:
-                # remove all trailing lines
-                self.client.reset_request()
-                break
+        try:
+            while True:
+                chunk_size = self.__get_chunk_size()
+                logging.debug("chunk-size: %s", chunk_size)
+                if chunk_size == 0:
+                    # remove all trailing lines
+                    self.client.reset_request()
+                    break
 
-            buffer = self.client.read(chunk_size)
-            yield buffer
+                buffer = self.client.read(chunk_size)
+                yield buffer
 
-            self.client.read_line()  # remove trailing CRLF
+                self.client.read_line()  # remove trailing CRLF
+
+        except TimeoutError:
+            raise IncompleteResponse("Timed out before receiving the complete payload!")
+        except ConnectionError:
+            raise IncompleteResponse("Connection closed before receiving the complete payload!")
 
     def __get_chunk_size(self):
         line = self.client.read_line()