Repository: scikit-hep/pyhf

Issue:
Make download tests robust to HEPData outages by mocking of tar.gz files
@kratsg is pointing out that we should probably have these mocked so that the test suite doesn't rely on HEPData behavior.

Example of something similar he's done recently

```python
def test_user_expires_reauthenticate(user_temp, requests_mock, mocker):
    user, _ = user_temp
    assert user.is_authenticated()
    assert user.is_expired() == False
    assert user.expires_in > 0
    user._id_token['exp'] = time.time() - 1
    assert user.is_authenticated()
    assert user.is_expired()
    assert user.expires_in == 0

    mock = mocker.patch.object(user, '_parse_id_token')
    requests_mock.post(
        requests.compat.urljoin(itkdb.settings.ITKDB_AUTH_URL, 'grantToken'),
        text=json.dumps(
            {
                'id_token': {
                    'exp': time.time() + 3600,
                    'name': _name,
                    'uuidentity': _identity,
                },
                'access_token': _access_token,
            }
        ),
    )
    user.authenticate()
    user._id_token = {'exp': time.time() + 3600, 'name': _name, 'uuidentity': _identity}
    assert user.is_authenticated()
    assert user.is_expired() == False
```

_Originally posted by @matthewfeickert in https://github.com/scikit-hep/pyhf/issues/1485#issuecomment-857370600_


Discussion:

Resolving patch:
diff --git a/setup.py b/setup.py
--- a/setup.py
+++ b/setup.py
@@ -32,6 +32,7 @@
             'pytest>=6.0',
             'pytest-cov>=2.5.1',
             'pytest-mock',
+            'requests-mock>=1.9.0',
             'pytest-benchmark[histogram]',
             'pytest-console-scripts',
             'pytest-mpl',
diff --git a/src/pyhf/contrib/utils.py b/src/pyhf/contrib/utils.py
--- a/src/pyhf/contrib/utils.py
+++ b/src/pyhf/contrib/utils.py
@@ -1,9 +1,13 @@
 """Helper utilities for common tasks."""
 
-from urllib.parse import urlparse
+import logging
 import tarfile
+import zipfile
 from io import BytesIO
-import logging
+from pathlib import Path
+from shutil import rmtree
+from urllib.parse import urlparse
+
 from pyhf import exceptions
 
 log = logging.getLogger(__name__)
@@ -62,17 +66,66 @@ def download(archive_url, output_directory, force=False, compress=False):
         # The HEPData landing page for the resource file can check if the Accept
         # request HTTP header matches the content type of the resource file and
         # return the content directly if so.
+        # TODO: Figure out how to accept headers of both application/x-tar and
+        # application/zip.
         with requests.get(
             archive_url, headers={"Accept": "application/x-tar"}
         ) as response:
+            if response.status_code != 200:
+                raise exceptions.InvalidArchive(
+                    f"{archive_url} gives a response code of {response.status_code}.\n"
+                    + "There is either something temporarily wrong with the archive host"
+                    + f" or {archive_url} is an invalid URL."
+                )
+
             if compress:
                 with open(output_directory, "wb") as archive:
                     archive.write(response.content)
             else:
-                with tarfile.open(
-                    mode="r|gz", fileobj=BytesIO(response.content)
-                ) as archive:
-                    archive.extractall(output_directory)
+                # Support for file-like objects for tarfile.is_tarfile was added
+                # in Python 3.9, so as pyhf is currently Python 3.7+ then can't
+                # do tarfile.is_tarfile(BytesIO(response.content)).
+                # Instead, just use a 'try except' block to determine if the
+                # archive is a valid tarfile.
+                # TODO: Simplify after pyhf is Python 3.9+ only
+                try:
+                    # Use transparent compression to allow for .tar or .tar.gz
+                    with tarfile.open(
+                        mode="r:*", fileobj=BytesIO(response.content)
+                    ) as archive:
+                        archive.extractall(output_directory)
+                except tarfile.ReadError:
+                    if not zipfile.is_zipfile(BytesIO(response.content)):
+                        raise exceptions.InvalidArchive(
+                            f"The archive downloaded from {archive_url} is not a tarfile"
+                            + " or a zipfile and so can not be opened as one."
+                        )
+
+                    output_directory = Path(output_directory)
+                    if output_directory.exists():
+                        rmtree(output_directory)
+                    with zipfile.ZipFile(BytesIO(response.content)) as archive:
+                        archive.extractall(output_directory)
+
+                        # zipfile.ZipFile.extractall extracts to a directory
+                        # below a target directory, so to match the extraction
+                        # path of tarfile.TarFile.extractall move the extracted
+                        # directory to a temporary path and then replace the
+                        # output directory target with the contents at the
+                        # temporary path.
+                        # The directory is moved instead of being extracted one
+                        # directory up and then renamed as the name of the
+                        # zipfile directory is set at zipfile creation time and
+                        # isn't knowable in advance.
+                        child_path = [child for child in output_directory.iterdir()][0]
+                        _tmp_path = output_directory.parent.joinpath(
+                            Path(output_directory.name + "__tmp__")
+                        )
+                        child_path.replace(_tmp_path)
+                        # the zipfile could contain remnant __MACOSX directories
+                        # from creation time
+                        rmtree(output_directory)
+                        _tmp_path.replace(output_directory)
 
 
 except ModuleNotFoundError:
diff --git a/src/pyhf/exceptions/__init__.py b/src/pyhf/exceptions/__init__.py
--- a/src/pyhf/exceptions/__init__.py
+++ b/src/pyhf/exceptions/__init__.py
@@ -70,6 +70,10 @@ class InvalidArchiveHost(Exception):
     """InvalidArchiveHost is raised when a given patchset archive url is not an approved host."""
 
 
+class InvalidArchive(Exception):
+    """InvalidArchive is raised when a given patchset archive url does not return a valid response."""
+
+
 class InvalidPatchSet(Exception):
     """InvalidPatchSet is raised when a given patchset object does not have the right configuration, even though it validates correctly against the schema."""
 


Test patch:
diff --git a/tests/contrib/test_utils.py b/tests/contrib/test_utils.py
new file mode 100644
--- /dev/null
+++ b/tests/contrib/test_utils.py
@@ -0,0 +1,99 @@
+import tarfile
+import zipfile
+from pathlib import Path
+from shutil import rmtree
+
+import pytest
+
+from pyhf.contrib.utils import download
+from pyhf.exceptions import InvalidArchive, InvalidArchiveHost
+
+
+@pytest.fixture(scope="function")
+def tarfile_path(tmpdir):
+    with open(tmpdir.join("test_file.txt").strpath, "w") as write_file:
+        write_file.write("test file")
+    with tarfile.open(tmpdir.join("test_tar.tar.gz").strpath, mode="w:gz") as archive:
+        archive.add(tmpdir.join("test_file.txt").strpath)
+    return Path(tmpdir.join("test_tar.tar.gz").strpath)
+
+
+@pytest.fixture(scope="function")
+def tarfile_uncompressed_path(tmpdir):
+    with open(tmpdir.join("test_file.txt").strpath, "w") as write_file:
+        write_file.write("test file")
+    with tarfile.open(tmpdir.join("test_tar.tar").strpath, mode="w") as archive:
+        archive.add(tmpdir.join("test_file.txt").strpath)
+    return Path(tmpdir.join("test_tar.tar").strpath)
+
+
+@pytest.fixture(scope="function")
+def zipfile_path(tmpdir):
+    with open(tmpdir.join("test_file.txt").strpath, "w") as write_file:
+        write_file.write("test file")
+    with zipfile.ZipFile(tmpdir.join("test_zip.zip").strpath, "w") as archive:
+        archive.write(tmpdir.join("test_file.txt").strpath)
+    return Path(tmpdir.join("test_zip.zip").strpath)
+
+
+def test_download_untrusted_archive_host(tmpdir, requests_mock):
+    archive_url = "https://www.pyhfthisdoesnotexist.org"
+    requests_mock.get(archive_url)
+
+    with pytest.raises(InvalidArchiveHost):
+        download(archive_url, tmpdir.join("likelihoods").strpath)
+
+
+def test_download_invalid_archive(tmpdir, requests_mock):
+    archive_url = "https://www.hepdata.net/record/resource/1408476?view=true"
+    requests_mock.get(archive_url, status_code=404)
+
+    with pytest.raises(InvalidArchive):
+        download(archive_url, tmpdir.join("likelihoods").strpath)
+
+
+def test_download_compress(tmpdir, requests_mock):
+    archive_url = "https://www.hepdata.net/record/resource/1408476?view=true"
+    requests_mock.get(archive_url)
+
+    download(archive_url, tmpdir.join("likelihoods").strpath, compress=True)
+
+
+def test_download_archive_type(
+    tmpdir, mocker, requests_mock, tarfile_path, tarfile_uncompressed_path, zipfile_path
+):
+    archive_url = "https://www.hepdata.net/record/resource/1408476?view=true"
+    output_directory = tmpdir.join("likelihoods").strpath
+    # Give BytesIO a tarfile
+    requests_mock.get(archive_url, content=open(tarfile_path, "rb").read())
+    download(archive_url, output_directory)
+
+    # Give BytesIO an uncompressed tarfile
+    requests_mock.get(archive_url, content=open(tarfile_uncompressed_path, "rb").read())
+    download(archive_url, output_directory)
+
+    # Give BytesIO a zipfile
+    requests_mock.get(archive_url, content=open(zipfile_path, "rb").read())
+    # Run without and with existing output_directory to cover both
+    # cases of the shutil.rmtree logic
+    rmtree(Path(output_directory))
+    download(archive_url, output_directory)  # without
+    download(archive_url, output_directory)  # with
+
+    # Give BytesIO a zipfile (using same requests_mock as previous) but have
+    # zipfile.is_zipfile reject it
+    mocker.patch("zipfile.is_zipfile", return_value=False)
+    with pytest.raises(InvalidArchive):
+        download(archive_url, output_directory)
+
+
+def test_download_archive_force(tmpdir, requests_mock, tarfile_path):
+    archive_url = "https://www.cern.ch/record/resource/123456789"
+    requests_mock.get(
+        archive_url, content=open(tarfile_path, "rb").read(), status_code=200
+    )
+
+    with pytest.raises(InvalidArchiveHost):
+        download(archive_url, tmpdir.join("likelihoods").strpath, force=False)
+
+    download(archive_url, tmpdir.join("likelihoods").strpath, force=True)
diff --git a/tests/test_scripts.py b/tests/test_scripts.py
--- a/tests/test_scripts.py
+++ b/tests/test_scripts.py
@@ -1,14 +1,26 @@
 import json
+import logging
 import shlex
-import pyhf
-import time
 import sys
-import logging
+import tarfile
+import time
+from importlib import import_module, reload
+from pathlib import Path
+from unittest import mock
+
 import pytest
 from click.testing import CliRunner
-from unittest import mock
-from importlib import reload
-from importlib import import_module
+
+import pyhf
+
+
+@pytest.fixture(scope="function")
+def tarfile_path(tmpdir):
+    with open(tmpdir.join("test_file.txt").strpath, "w") as write_file:
+        write_file.write("test file")
+    with tarfile.open(tmpdir.join("test_tar.tar.gz").strpath, mode="w:gz") as archive:
+        archive.add(tmpdir.join("test_file.txt").strpath)
+    return Path(tmpdir.join("test_tar.tar.gz").strpath)
 
 
 def test_version(script_runner):
@@ -273,9 +285,10 @@ def test_testpoi(tmpdir, script_runner):
         results_exp.append(d['CLs_exp'])
         results_obs.append(d['CLs_obs'])
 
-    import numpy as np
     import itertools
 
+    import numpy as np
+
     for pair in itertools.combinations(results_exp, r=2):
         assert not np.array_equal(*pair)
 
@@ -543,7 +556,8 @@ def test_workspace_digest(tmpdir, script_runner, algorithms, do_json):
         "https://doi.org/10.17182/hepdata.89408.v1/r2",
     ],
 )
-def test_patchset_download(tmpdir, script_runner, archive):
+def test_patchset_download(tmpdir, script_runner, requests_mock, tarfile_path, archive):
+    requests_mock.get(archive, content=open(tarfile_path, "rb").read())
     command = f'pyhf contrib download {archive} {tmpdir.join("likelihoods").strpath}'
     ret = script_runner.run(*shlex.split(command))
     assert ret.success
@@ -553,6 +567,9 @@ def test_patchset_download(tmpdir, script_runner, archive):
     ret = script_runner.run(*shlex.split(command))
     assert ret.success
 
+    requests_mock.get(
+        "https://www.pyhfthisdoesnotexist.org/record/resource/1234567", status_code=200
+    )
     command = f'pyhf contrib download --verbose https://www.pyhfthisdoesnotexist.org/record/resource/1234567 {tmpdir.join("likelihoods").strpath}'
     ret = script_runner.run(*shlex.split(command))
     assert not ret.success
@@ -560,11 +577,15 @@ def test_patchset_download(tmpdir, script_runner, archive):
         "pyhf.exceptions.InvalidArchiveHost: www.pyhfthisdoesnotexist.org is not an approved archive host"
         in ret.stderr
     )
-    # Force a download from a real URL, but one that doesn't have an existing file
+
+    # httpstat.us is a real wesite that can be used for testing responses
+    requests_mock.get(
+        "https://httpstat.us/404/record/resource/1234567", status_code=404
+    )
     command = f'pyhf contrib download --verbose --force https://httpstat.us/404/record/resource/1234567 {tmpdir.join("likelihoods").strpath}'
     ret = script_runner.run(*shlex.split(command))
     assert not ret.success
-    assert "tarfile.ReadError: not a gzip file" in ret.stderr
+    assert "gives a response code of 404" in ret.stderr
 
 
 def test_missing_contrib_extra(caplog):