diff --git a/src/tufup/client.py b/src/tufup/client.py index aa728b1..8607a7e 100644 --- a/src/tufup/client.py +++ b/src/tufup/client.py @@ -8,11 +8,8 @@ from urllib import parse import requests -from requests.adapters import ReadTimeoutError from requests.auth import AuthBase -from tuf.api.exceptions import ( - DownloadError, SlowRetrievalError, UnsignedMetadataError -) +from tuf.api.exceptions import DownloadError, UnsignedMetadataError from tuf.api.metadata import TargetFile import tuf.ngclient # RequestsFetcher is "private", but we'll just have to live with that, for now. @@ -345,23 +342,7 @@ def _get_session(self, url: str) -> requests.Session: return session def _chunks(self, response: "requests.Response") -> Iterator[bytes]: - """ - Override _chunks() to: - - prevent automatic decoding of gzip files (python-tuf issue #2047) - - call progress hook - - todo: adapt, if necessary, when a fix for python-tuf #2047 is released - """ - try: - while True: - data = response.raw.read( - amt=self.chunk_size, decode_content=False - ) - if not data: - break - self._progress(bytes_new=len(data)) - yield data - except ReadTimeoutError as e: - raise SlowRetrievalError from e - finally: - response.close() + """Call progress hook for every chunk.""" + for data in super()._chunks(response=response): + self._progress(bytes_new=len(data)) + yield data diff --git a/tests/test_client.py b/tests/test_client.py index f63fad9..3cb3f05 100644 --- a/tests/test_client.py +++ b/tests/test_client.py @@ -278,11 +278,10 @@ def test__chunks_without_progress_hook(self): chunk_count = 10 chunks = [b'x' * chunk_size] * chunk_count - def mock_read(**kwargs): - if chunks: - return chunks.pop() + def mock_iter_content(*args): + yield from chunks - mock_response = Mock(raw=Mock(read=mock_read), close=Mock()) + mock_response = Mock(iter_content=mock_iter_content, close=Mock()) fetcher = AuthRequestsFetcher() fetcher.chunk_size = chunk_size # _chunks should work even if attach_progress_hook was not called @@ -297,11 +296,10 @@ def test__chunks_with_progress_hook(self): chunk_count = 10 chunks = [b'x' * chunk_size] * chunk_count - def mock_read(**kwargs): - if chunks: - return chunks.pop() + def mock_iter_content(*args): + yield from chunks - mock_response = Mock(raw=Mock(read=mock_read), close=Mock()) + mock_response = Mock(iter_content=mock_iter_content, close=Mock()) fetcher = AuthRequestsFetcher() fetcher.chunk_size = chunk_size # test custom progress hook