Take out connection time from download speed calculation

- Update RequestsFetcher.fetch to return a generator object.
- Update _download_file 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

@@ -190,33 +190,33 @@ def _download_file(url, required_length, fetcher, 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 fetcher.fetch(url, required_length):
+    chunks = 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?
     _check_downloaded_length(number_of_bytes_received, required_length,
                              STRICT_REQUIRED_LENGTH=STRICT_REQUIRED_LENGTH,

tuf/fetcher.py

@@ -96,12 +96,12 @@ 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=tuf.settings.SOCKET_TIMEOUT) as response:
-
-      # Check response status.
-      response.raise_for_status()
+    response = session.get(url, stream=True,
+        timeout=tuf.settings.SOCKET_TIMEOUT)
+    # Check response status.
+    response.raise_for_status()
 
+    def chunks():
       try:
         bytes_received = 0
         while True:
@@ -122,11 +122,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) + '/' +
@@ -135,9 +130,18 @@ 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))
 
+      finally:
+        response.close()
+
+    return chunks()