Take out connection time from download speed calculation

- Update RequestsFetcher.fetch to return a generator object.
- Update FileDownload to skip connection time when calculating
average download speed.
- Write chunk to temp file before exiting the fetcher loop
on error.

Signed-off-by: Teodora Sechkova <[email protected]>

Author: sechkova

tuf/download.py

@@ -204,33 +204,35 @@ def _download_file(self, url, required_length, STRICT_REQUIRED_LENGTH=True):
     # This is the temporary file that we will return to contain the contents of
     # the downloaded file.
     temp_file = tempfile.TemporaryFile()
-    start_time = timeit.default_timer()
 
     average_download_speed = 0
     number_of_bytes_received = 0
 
     try:
 
-      for chunk in self.fetcher.fetch(url, required_length):
+      chunks = self.fetcher.fetch(url, required_length)
+      start_time = timeit.default_timer()
+      for chunk in chunks:
 
         stop_time = timeit.default_timer()
-        seconds_spent_receiving = stop_time - start_time
+        temp_file.write(chunk)
+
         # Measure the average download speed.
         number_of_bytes_received += len(chunk)
+        seconds_spent_receiving = stop_time - start_time
         average_download_speed = number_of_bytes_received / seconds_spent_receiving
 
         if average_download_speed < tuf.settings.MIN_AVERAGE_DOWNLOAD_SPEED:
           logger.debug('The average download speed dropped below the minimum'
-            ' average download speed set in tuf.settings.py.')
+            ' average download speed set in tuf.settings.py. Stopping the'
+            ' download!')
           break
 
         else:
           logger.debug('The average download speed has not dipped below the'
             ' minimum average download speed set in tuf.settings.py.')
 
 
-        temp_file.write(chunk)
-
       # Does the total number of downloaded bytes match the required length?
       self._check_downloaded_length(number_of_bytes_received, required_length,
                                STRICT_REQUIRED_LENGTH=STRICT_REQUIRED_LENGTH,

tuf/fetcher.py

@@ -113,12 +113,11 @@ def fetch(self, url, required_length):
     # Always set the timeout. This timeout value is interpreted by requests as:
     #  - connect timeout (max delay before first byte is received)
     #  - read (gap) timeout (max delay between bytes received)
-    with session.get(url, stream=True,
-        timeout=self.SOCKET_TIMEOUT) as response:
-
-      # Check response status.
-      response.raise_for_status()
+    response = session.get(url, stream=True, timeout=self.SOCKET_TIMEOUT)
+    # Check response status.
+    response.raise_for_status()
 
+    def chunks():
       try:
         bytes_received = 0
         while True:
@@ -139,11 +138,6 @@ def fetch(self, url, required_length):
           data = response.raw.read(read_amount)
           bytes_received += len(data)
 
-          yield data
-
-          if bytes_received == required_length:
-            break
-
           # We might have no more data to read. Check number of bytes downloaded.
           if not data:
             logger.debug('Downloaded ' + repr(bytes_received) + '/' +
@@ -152,9 +146,15 @@ def fetch(self, url, required_length):
             # Finally, we signal that the download is complete.
             break
 
+          yield data
+
+          if bytes_received == required_length:
+            break
+
       except urllib3.exceptions.ReadTimeoutError as e:
         raise tuf.exceptions.SlowRetrievalError(str(e))
 
+    return chunks()