robust download (#100)

* implement robust download

* robust download

* better tests

* cleanup

Author: malmans2

cads_api_client/processing.py

@@ -678,6 +678,18 @@ def asset(self) -> dict[str, Any]:
         """
         return dict(self.json["asset"]["value"])
 
+    def _download(self, url: str, target: str) -> requests.Response:
+        download_options = {"stream": True, "resume_transfers": True}
+        download_options.update(self.download_options)
+        multiurl.download(
+            url,
+            target=target,
+            **self.retry_options,
+            **self.request_options,
+            **download_options,
+        )
+        return requests.Response()  # mutliurl robust needs a response
+
     def download(
         self,
         target: str | None = None,
@@ -699,15 +711,11 @@ def download(
             parts = urllib.parse.urlparse(url)
             target = parts.path.strip("/").split("/")[-1]
 
-        download_options = {"stream": True}
-        download_options.update(self.download_options)
-        multiurl.download(
-            url,
-            target=target,
-            **self.retry_options,
-            **self.request_options,
-            **download_options,
-        )
+        if os.path.exists(target):
+            os.remove(target)
+
+        robust_download = multiurl.robust(self._download, **self.retry_options)
+        robust_download(url, target)
         self._check_size(target)
         return target
 

tests/integration_test_40_results.py

@@ -1,10 +1,16 @@
+import contextlib
 import os
 import pathlib
+import random
+from typing import Any
 
 import pytest
+import requests
 
 from cads_api_client import ApiClient, Results
 
+does_not_raise = contextlib.nullcontext
+
 
 @pytest.fixture
 def results(api_anon_client: ApiClient) -> Results:
@@ -45,3 +51,60 @@ def test_results_progress(
     submitted.download(target=str(tmp_path / "test.grib"))
     captured = capsys.readouterr()
     assert captured.err if progress else not captured.err
+
+
+@pytest.mark.parametrize(
+    "maximum_tries,raises",
+    [
+        (500, does_not_raise()),
+        (1, pytest.raises(requests.ConnectionError, match="Random error.")),
+    ],
+)
+def test_results_robust_download(
+    api_root_url: str,
+    api_anon_key: str,
+    monkeypatch: pytest.MonkeyPatch,
+    tmp_path: pathlib.Path,
+    maximum_tries: int,
+    raises: contextlib.nullcontext[Any],
+) -> None:
+    from multiurl.http import FullHTTPDownloader
+
+    def patched_iter_content(self, *args, **kwargs):  # type: ignore
+        for chunk in self.iter_content(chunk_size=1):
+            if random.choice([True, False]):
+                raise requests.ConnectionError("Random error.")
+            yield chunk
+
+    def make_stream(self):  # type: ignore
+        request = self.issue_request(self.range)
+        return request.patched_iter_content
+
+    client = ApiClient(
+        url=api_root_url, key=api_anon_key, retry_after=0, maximum_tries=maximum_tries
+    )
+    results = client.submit_and_wait_on_results("test-adaptor-dummy", size=10)
+    monkeypatch.setattr(
+        requests.Response, "patched_iter_content", patched_iter_content, raising=False
+    )
+    monkeypatch.setattr(FullHTTPDownloader, "make_stream", make_stream)
+
+    target = tmp_path / "test.grib"
+    with raises:
+        results.download(str(target))
+
+
+def test_results_override(api_anon_client: ApiClient, tmp_path: pathlib.Path) -> None:
+    target_1 = tmp_path / "tmp1.grib"
+    api_anon_client.retrieve("test-adaptor-dummy", size=1, target=str(target_1))
+
+    target_2 = tmp_path / "tmp2.grib"
+    api_anon_client.retrieve("test-adaptor-dummy", size=2, target=str(target_2))
+
+    target = tmp_path / "tmp.grib"
+    api_anon_client.retrieve("test-adaptor-dummy", size=1, target=str(target))
+    assert target.read_bytes() == target_1.read_bytes()
+    api_anon_client.retrieve("test-adaptor-dummy", size=2, target=str(target))
+    assert target.read_bytes() == target_2.read_bytes()
+    api_anon_client.retrieve("test-adaptor-dummy", size=1, target=str(target))
+    assert target.read_bytes() == target_1.read_bytes()