From 91e93d0b4dad1e4df9605f9432eb876ac6fc0ff3 Mon Sep 17 00:00:00 2001
From: Mattia Almansi <m.almansi@bopen.eu>
Date: Wed, 25 Sep 2024 19:36:31 +0200
Subject: [PATCH] full support for download

---
 cads_api_client/legacy_api_client.py          | 19 ++++++++++++++++++-
 .../integration_test_70_legacy_api_client.py  | 16 ++++++++++++----
 2 files changed, 30 insertions(+), 5 deletions(-)

diff --git a/cads_api_client/legacy_api_client.py b/cads_api_client/legacy_api_client.py
index 1eb9283..e20a519 100644
--- a/cads_api_client/legacy_api_client.py
+++ b/cads_api_client/legacy_api_client.py
@@ -220,7 +220,24 @@ def _download(self, results, targets=None):
             return [self._download(x, targets) for x in results]
 
         if isinstance(results, dict):
-            self.raise_not_implemented_error()
+            if "location" in results and "contentLength" in results:
+                reply = dict(
+                    location=results["location"],
+                    content_length=results["contentLength"],
+                    content_type=results.get("contentType"),
+                )
+
+                if targets:
+                    path = targets.pop(0)
+                else:
+                    path = None
+
+                return cdsapi.api.Result(self, reply).download(path)
+
+            r = {}
+            for v in results.values():
+                r[v] = self._download(v, targets)
+            return r
 
         return results
 
diff --git a/tests/integration_test_70_legacy_api_client.py b/tests/integration_test_70_legacy_api_client.py
index 099ecd6..d74f2d1 100644
--- a/tests/integration_test_70_legacy_api_client.py
+++ b/tests/integration_test_70_legacy_api_client.py
@@ -226,7 +226,7 @@ def test_legacy_api_client_download(
     client = LegacyApiClient(
         url=api_root_url,
         key=api_anon_key,
-        retry_max=0,
+        retry_max=1,
         wait_until_complete=False,
     )
     remote = client.retrieve("test-adaptor-dummy", {"size": 1})
@@ -234,11 +234,19 @@ def test_legacy_api_client_download(
     target = client.download(remote)
     assert os.path.getsize(target) == 1
 
-    results = (remote, remote.make_results())
+    results = remote.make_results()
+    results_dict = {
+        "location": results.location,
+        "contentLength": results.content_length,
+    }
+    results_tuple = (remote, remote.make_results(), results_dict)
     target1 = "remote.grib"
     target2 = "results.grib"
-    assert client.download(results, [target1, target2]) == [target1, target2]
-    assert os.path.getsize(target1) == os.path.getsize(target2) == 1
+    target3 = "dict.grib"
+    targets = [target1, target2, target3]
+    assert client.download(results_tuple, targets) == [target1, target2, target3]
+    assert all(os.path.getsize(target) == 1 for target in targets)
+    os.path.getsize(target1) == os.path.getsize(target2) == 1
 
 
 def test_legacy_api_client_status(legacy_client: LegacyApiClient) -> None: