sandorkertesz · Mar 13, 2024 · Mar 6, 2024 · Mar 6, 2024 · Mar 11, 2024 · Mar 11, 2024
diff --git a/multiurl/base.py b/multiurl/base.py
@@ -71,6 +71,7 @@ def __init__(
         resume_transfers=False,
         override_target_file=True,
         download_file_extension=None,
+        auth=None,
         **kwargs,
     ):
         self.url = url
@@ -83,6 +84,7 @@ def __init__(
         self.resume_transfers = resume_transfers
         self.override_target_file = override_target_file
         self.download_file_extension = download_file_extension
+        self.auth = auth
 
     def mutate(self, *args, **kwargs):
         return self

diff --git a/multiurl/http.py b/multiurl/http.py
@@ -75,6 +75,7 @@ def headers(self):
                         verify=self.verify,
                         timeout=self.timeout,
                         allow_redirects=True,
+                        auth=self.auth,
                     )
                     r.raise_for_status()
                     for k, v in r.headers.items():
@@ -222,6 +223,7 @@ def issue_request(self, bytes_ranges=None):
             verify=self.verify,
             timeout=self.timeout,
             headers=headers,
+            auth=self.auth,
         )
         try:
             r.raise_for_status()

diff --git a/tests/test_auth.py b/tests/test_auth.py
@@ -0,0 +1,114 @@
+# (C) Copyright 2021 ECMWF.
+#
+# This software is licensed under the terms of the Apache Licence Version 2.0
+# which can be obtained at http://www.apache.org/licenses/LICENSE-2.0.
+# In applying this licence, ECMWF does not waive the privileges and immunities
+# granted to it by virtue of its status as an intergovernmental organisation
+# nor does it submit to any jurisdiction.
+#
+
+import logging
+import os
+
+from multiurl import download
+
+# NOTE: we just test if the auth object is properly called with the
+# requests when using download()
+
+
+class Auth:
+    def __init__(self):
+        from collections import defaultdict
+
+        self.calls = defaultdict(set)
+
+    def __call__(self, r):
+        method = r.method.lower()
+        self.calls[method].add(r.url)
+        return r
+
+
+def test_auth_single():
+    auth = Auth()
+    url = "http://get.ecmwf.int/test-data/metview/gallery/temp.bufr"
+    download(url=url, target="out.data", auth=auth)
+
+    assert auth.calls["head"] == set([url])
+    assert auth.calls["get"] == set([url])
+
+
+def test_auth_single_fake_headers():
+    auth = Auth()
+    url = "http://get.ecmwf.int/test-data/metview/gallery/temp.bufr"
+    download(url=url, target="out.data", auth=auth, fake_headers={})
+
+    assert auth.calls["head"] == set()
+    assert auth.calls["get"] == set([url])
+
+
+def test_auth_single_parts():
+    auth = Auth()
+    url = "http://get.ecmwf.int/test-data/metview/gallery/temp.bufr"
+
+    download(url=url, target="out.data", parts=((0, 4),), auth=auth)
+
+    assert auth.calls["head"] == set([url])
+    assert auth.calls["get"] == set([url])
+    assert os.path.getsize("out.data") == 4
+
+    with open("out.data", "rb") as f:
+        assert f.read() == b"BUFR"
+
+
+def test_auth_single_multi_parts():
+    auth = Auth()
+    url = "http://get.ecmwf.int/test-data/metview/gallery/temp.bufr"
+
+    download(url=url, target="out.data", parts=((0, 4), (20, 4)), auth=auth)
+
+    assert auth.calls["head"] == set([url])
+    assert auth.calls["get"] == set([url])
+    assert os.path.getsize("out.data") == 8
+
+    with open("out.data", "rb") as f:
+        assert f.read(4) == b"BUFR"
+
+
+def test_auth_multi():
+    auth = Auth()
+    urls = [
+        "http://get.ecmwf.int/test-data/earthkit-data/examples/test.grib",
+        "http://get.ecmwf.int/test-data/earthkit-data/examples/test6.grib",
+    ]
+
+    download(url=urls, target="out.data", auth=auth)
+
+    assert auth.calls["head"] == set(urls)
+    assert auth.calls["get"] == set(urls)
+    assert os.path.getsize("out.data") == 2492
+
+    with open("out.data", "rb") as f:
+        assert f.read(4) == b"GRIB"
+
+
+def test_auth_multi_parts():
+    auth = Auth()
+    urls = [
+        "http://get.ecmwf.int/test-data/earthkit-data/examples/test.grib",
+        "http://get.ecmwf.int/test-data/earthkit-data/examples/test6.grib",
+    ]
+
+    download(url=urls, target="out.data", parts=((0, 4),), auth=auth)
+
+    assert auth.calls["head"] == set(urls)
+    assert auth.calls["get"] == set(urls)
+    assert os.path.getsize("out.data") == 8
+
+    with open("out.data", "rb") as f:
+        assert f.read(4) == b"GRIB"
+        assert f.read(4) == b"GRIB"
+
+
+if __name__ == "__main__":
+    logging.basicConfig(level=logging.DEBUG)
+    # test_order()