Test downloading data in more than one chunk

sechkova · sechkova · commit 50b3b193925d · 2021-02-24T11:41:22.000+02:00
Add test cases to test_fetcher and test_download that
decrease default chunk size and download data in more
than one chunk.

Small code-style improvements.

Signed-off-by: Teodora Sechkova &lt;tsechkova@vmware.com&gt;
diff --git a/tests/test_download.py b/tests/test_download.py
@@ -112,6 +112,29 @@ def test_download_url_to_tempfileobj(self):
       self.assertEqual(self.target_data_length, len(temp_file_data))
 
 
+  # Test: Download url in more than one chunk.
+  def test_download_url_in_chunks(self):
+
+    # Set smaller chunk size to ensure that the file will be downloaded
+    # in more than one chunk
+    default_chunk_size = tuf.settings.CHUNK_SIZE
+    tuf.settings.CHUNK_SIZE = 4
+    # We don't have access to chunks from  download_file()
+    # so we just confirm that the expectation of more than one chunk is
+    # correct and verify that no errors are raised during download
+    chunks_count = self.target_data_length/tuf.settings.CHUNK_SIZE
+    self.assertGreater(chunks_count, 1)
+
+    download_file = download.safe_download
+    with download_file(self.url, self.target_data_length, self.fetcher) as temp_fileobj:
+      temp_fileobj.seek(0)
+      temp_file_data = temp_fileobj.read().decode('utf-8')
+      self.assertEqual(self.target_data, temp_file_data)
+      self.assertEqual(self.target_data_length, len(temp_file_data))
+
+    # Restore default settings
+    tuf.settings.CHUNK_SIZE = default_chunk_size
+
 
   # Test: Incorrect lengths.
   def test_download_url_to_tempfileobj_and_lengths(self):
diff --git a/tests/test_fetcher.py b/tests/test_fetcher.py
@@ -16,13 +16,17 @@
 <Purpose>
   Unit test for RequestsFetcher.
 """
+# Help with Python 2 compatibility, where the '/' operator performs
+# integer division.
+from __future__ import division
 
 import logging
 import os
 import io
 import sys
 import unittest
 import tempfile
+import math
 
 import tuf
 import tuf.exceptions
@@ -31,7 +35,6 @@
 
 from tests import utils
 
-
 logger = logging.getLogger(__name__)
 
 
@@ -77,15 +80,12 @@ def tearDown(self):
 
   # Test: Normal case.
   def test_fetch(self):
-
     for chunk in self.fetcher.fetch(self.url, self.file_length):
         self.temp_file.write(chunk)
 
     self.temp_file.seek(0)
     temp_file_data = self.temp_file.read().decode('utf-8')
     self.assertEqual(self.file_contents, temp_file_data)
-    self.assertEqual(self.file_length, len(temp_file_data))
-
 
   # Test if fetcher downloads file up to a required length
   def test_fetch_restricted_length(self):
@@ -96,7 +96,7 @@ def test_fetch_restricted_length(self):
     self.assertEqual(self.temp_file.tell(), self.file_length-4)
 
 
-  # Test if fetcher does not downlad more than actual file length
+  # Test that fetcher does not download more than actual file length
   def test_fetch_upper_length(self):
     for chunk in self.fetcher.fetch(self.url, self.file_length+4):
         self.temp_file.write(chunk)
@@ -107,9 +107,34 @@ def test_fetch_upper_length(self):
 
   # Test incorrect URL parsing
   def test_url_parsing(self):
-   with self.assertRaises(tuf.exceptions.URLParsingError) as cm:
-     for chunk in self.fetcher.fetch(self.random_string(), self.file_length):
-         self.temp_file.write(chunk)
+    with self.assertRaises(tuf.exceptions.URLParsingError):
+      self.fetcher.fetch(self.random_string(), self.file_length)
+
+
+  # Test: Normal case with url data downloaded in more than one chunk
+  def test_fetch_in_chunks(self):
+    # Set smaller chunk size to ensure that the file will be downloaded
+    # in more than one chunk
+    default_chunk_size = tuf.settings.CHUNK_SIZE
+    tuf.settings.CHUNK_SIZE = 4
+
+    # expected_chunks_count: 3
+    expected_chunks_count = math.ceil(self.file_length/tuf.settings.CHUNK_SIZE)
+    self.assertEqual(expected_chunks_count, 3)
+
+    chunks_count = 0
+    for chunk in self.fetcher.fetch(self.url, self.file_length):
+      self.temp_file.write(chunk)
+      chunks_count+=1
+
+    self.temp_file.seek(0)
+    temp_file_data = self.temp_file.read().decode('utf-8')
+    self.assertEqual(self.file_contents, temp_file_data)
+    # Check that we calculate chunks as expected
+    self.assertEqual(chunks_count, expected_chunks_count)
+
+    # Restore default settings
+    tuf.settings.CHUNK_SIZE = default_chunk_size
 
 
 
