From 15a1666f0888a76de4cffb51d6f29a49116f7fee Mon Sep 17 00:00:00 2001 From: Jon Palmer Date: Thu, 16 Jul 2026 13:19:17 -0700 Subject: [PATCH] fix: implement resilient HTTP Range resume download logic in utilities.py --- funannotate2/utilities.py | 61 ++++++++++++++++++++++++-- tests/unit/test_utilities.py | 85 ++++++++++++++++++++++++++++++++++++ 2 files changed, 143 insertions(+), 3 deletions(-) diff --git a/funannotate2/utilities.py b/funannotate2/utilities.py index 4b15f0c..4014ab2 100755 --- a/funannotate2/utilities.py +++ b/funannotate2/utilities.py @@ -242,10 +242,62 @@ def download(url, name, wget=False, timeout=60, retries=3): # Try HTTPS first while attempt < retries: try: + headers = {} + existing_size = 0 + if os.path.exists(file_name): + existing_size = os.path.getsize(file_name) + + # Send HEAD request to check remote size and Range support + remote_size = None + try: + # Use a smaller timeout for HEAD to avoid hanging + with requests.head(url, timeout=10, allow_redirects=True, verify=False) as head_resp: + if head_resp.status_code == 200: + remote_size = int(head_resp.headers.get("Content-Length", 0)) + except Exception: + pass + + # If local file matches remote size, download is complete + if remote_size and existing_size == remote_size: + return + + # If local file is larger than remote size, overwrite from scratch + if remote_size and existing_size > remote_size: + try: + os.remove(file_name) + except Exception: + pass + existing_size = 0 + + # If we have a partially downloaded file, try range request + if existing_size > 0: + headers["Range"] = f"bytes={existing_size}-" + # Use requests for HTTP/HTTPS - with requests.get(url, stream=True, timeout=timeout, verify=False) as r: + with requests.get(url, stream=True, timeout=timeout, verify=False, headers=headers) as r: + if r.status_code == 206: + write_mode = "ab" + elif r.status_code == 416: + # Range Not Satisfiable (might be fully downloaded or server error) + # Delete the file and start from scratch to be safe + try: + os.remove(file_name) + except Exception: + pass + existing_size = 0 + headers.pop("Range", None) + with requests.get(url, stream=True, timeout=timeout, verify=False) as r_retry: + r_retry.raise_for_status() + with open(file_name, "wb") as f: + for chunk in r_retry.iter_content(chunk_size=8192): + if chunk: + f.write(chunk) + return + else: + write_mode = "wb" + r.raise_for_status() # Raise an exception for HTTP errors - with open(file_name, "wb") as f: + with open(file_name, write_mode) as f: for chunk in r.iter_content(chunk_size=8192): if chunk: # filter out keep-alive chunks f.write(chunk) @@ -270,7 +322,10 @@ def download(url, name, wget=False, timeout=60, retries=3): return # If we get here, both HTTPS and FTP failed - raise Exception(f"Download failed: {str(e)}") + raise Exception( + f"Download failed: {str(e)}. If you have unstable internet or download timeouts, " + "try running with the '--wget' option to use wget for robust downloads." + ) def _download_ftp(url, file_name, timeout=60): diff --git a/tests/unit/test_utilities.py b/tests/unit/test_utilities.py index 57e2a52..717b3f3 100644 --- a/tests/unit/test_utilities.py +++ b/tests/unit/test_utilities.py @@ -8,6 +8,7 @@ from funannotate2.utilities import ( create_tmpdir, + download, merge_coordinates, naming_slug, readBlocks, @@ -259,3 +260,87 @@ def test_monitoring_does_not_debug_log_formatted_memory_report( debug_messages = [call.args[0] for call in logger.debug.call_args_list] assert "snap input.fasta" in debug_messages assert not any("Memory usage for" in message for message in debug_messages) + + +class TestDownload: + """Tests for the download function.""" + + @patch("funannotate2.utilities.requests.get") + @patch("funannotate2.utilities.requests.head") + def test_download_from_scratch(self, mock_head, mock_get, tmp_path): + """Test downloading a file from scratch.""" + dest = tmp_path / "test_file.txt" + + # Mock requests.head + mock_head_resp = MagicMock() + mock_head_resp.status_code = 200 + mock_head_resp.headers = {"Content-Length": "10"} + mock_head.return_value.__enter__.return_value = mock_head_resp + + # Mock requests.get + mock_get_resp = MagicMock() + mock_get_resp.status_code = 200 + mock_get_resp.iter_content.return_value = [b"0123456789"] + mock_get.return_value.__enter__.return_value = mock_get_resp + + download("https://example.com/file.txt", str(dest)) + + assert dest.exists() + assert dest.read_bytes() == b"0123456789" + mock_get.assert_called_once_with( + "https://example.com/file.txt", + stream=True, + timeout=60, + verify=False, + headers={}, + ) + + @patch("funannotate2.utilities.requests.get") + @patch("funannotate2.utilities.requests.head") + def test_download_resume(self, mock_head, mock_get, tmp_path): + """Test resuming a partial download using range requests.""" + dest = tmp_path / "test_file.txt" + dest.write_bytes(b"01234") # Partial file of 5 bytes + + # Mock requests.head + mock_head_resp = MagicMock() + mock_head_resp.status_code = 200 + mock_head_resp.headers = {"Content-Length": "10"} + mock_head.return_value.__enter__.return_value = mock_head_resp + + # Mock requests.get (returns range) + mock_get_resp = MagicMock() + mock_get_resp.status_code = 206 + mock_get_resp.iter_content.return_value = [b"56789"] + mock_get.return_value.__enter__.return_value = mock_get_resp + + download("https://example.com/file.txt", str(dest)) + + assert dest.exists() + assert dest.read_bytes() == b"0123456789" + mock_get.assert_called_once_with( + "https://example.com/file.txt", + stream=True, + timeout=60, + verify=False, + headers={"Range": "bytes=5-"}, + ) + + @patch("funannotate2.utilities.requests.get") + @patch("funannotate2.utilities.requests.head") + def test_download_already_complete(self, mock_head, mock_get, tmp_path): + """Test that if the file is already complete, no download request is made.""" + dest = tmp_path / "test_file.txt" + dest.write_bytes(b"0123456789") # 10 bytes + + # Mock requests.head + mock_head_resp = MagicMock() + mock_head_resp.status_code = 200 + mock_head_resp.headers = {"Content-Length": "10"} + mock_head.return_value.__enter__.return_value = mock_head_resp + + download("https://example.com/file.txt", str(dest)) + + # Check that get was never called + mock_get.assert_not_called() + assert dest.read_bytes() == b"0123456789"