Allow baseline_dir to be comma-separated URL list to allow mirrors to be specified

Author: astrofrog

.travis.yml

@@ -36,7 +36,7 @@ install:
 
 script:
    - python -c 'import pytest_mpl.plugin'
-   - py.test --mpl --cov pytest_mpl tests
+   - pytest -vv --mpl --cov pytest_mpl tests
    - python setup.py check --restructuredtext
 
 after_success:

CHANGES.md

@@ -1,7 +1,8 @@
 0.9 (unreleased)
 ----------------
 
-- No changes yet.
+- Allow baseline_dir to be comma-separated URL list to allow mirrors to
+  be specified. [#59]
 
 0.8 (2017-07-19)
 ----------------

README.rst

@@ -140,7 +140,9 @@ filename of the plot (which defaults to the name of the test with a
 The baseline directory in the decorator above will be interpreted as
 being relative to the test file. Note that the baseline directory can
 also be a URL (which should start with ``http://`` or ``https://`` and
-end in a slash).
+end in a slash). If you want to specify mirrors, set ``baseline_dir`` to
+a comma-separated list of URLs (real commas in the URL should be encoded
+as ``%2C``).
 
 Finally, you can also set a custom baseline directory globally when
 running tests by running ``py.test`` with:

pytest_mpl/plugin.py

@@ -42,16 +42,30 @@
 
 if sys.version_info[0] == 2:
     from urllib import urlopen
+    string_types = basestring
 else:
     from urllib.request import urlopen
-
-
-def _download_file(url):
-    u = urlopen(url)
+    string_types = str
+
+
+def _download_file(baseline, filename):
+    # Note that baseline can be a comma-separated list of URLs that we can
+    # then treat as mirrors
+    for base_url in baseline.split(','):
+        try:
+            u = urlopen(base_url + filename)
+            content = u.read()
+        except Exception as exc:
+            warnings.warn('Downloading {0} failed'.format(base_url + filename))
+        else:
+            break
+    else:
+        raise Exception("Could not download baseline image from any of the "
+                        "available URLs")
     result_dir = tempfile.mkdtemp()
     filename = os.path.join(result_dir, 'downloaded')
     with open(filename, 'wb') as tmpfile:
-        tmpfile.write(u.read())
+        tmpfile.write(content)
     return filename
 
 
@@ -60,9 +74,13 @@ def pytest_addoption(parser):
     group.addoption('--mpl', action='store_true',
                     help="Enable comparison of matplotlib figures to reference files")
     group.addoption('--mpl-generate-path',
-                    help="directory to generate reference images in, relative to location where py.test is run", action='store')
+                    help="directory to generate reference images in, relative "
+                    "to location where py.test is run", action='store')
     group.addoption('--mpl-baseline-path',
-                    help="directory containing baseline images, relative to location where py.test is run", action='store')
+                    help="directory containing baseline images, relative to "
+                    "location where py.test is run. This can also be a URL or a "
+                    "set of comma-separated URLs (in case mirrors are "
+                    "specified)", action='store')
 
     results_path_help = "directory for test results, relative to location where py.test is run"
     group.addoption('--mpl-results-path', help=results_path_help, action='store')
@@ -157,12 +175,12 @@ def item_function_wrapper(*args, **kwargs):
                     baseline_dir = os.path.join(os.path.dirname(item.fspath.strpath), 'baseline')
                 else:
                     baseline_dir = self.baseline_dir
+                baseline_remote = False
             else:
-                if not baseline_dir.startswith(('http://', 'https://')):
+                baseline_remote = baseline_dir.startswith(('http://', 'https://'))
+                if not baseline_remote:
                     baseline_dir = os.path.join(os.path.dirname(item.fspath.strpath), baseline_dir)
 
-            baseline_remote = baseline_dir.startswith('http')
-
             with plt.style.context(style), switch_backend(backend):
 
                 # Run test and get figure object
@@ -201,7 +219,7 @@ def item_function_wrapper(*args, **kwargs):
 
                     # Find path to baseline image
                     if baseline_remote:
-                        baseline_image_ref = _download_file(baseline_dir + filename)
+                        baseline_image_ref = _download_file(baseline_dir, filename)
                     else:
                         baseline_image_ref = os.path.abspath(os.path.join(os.path.dirname(item.fspath.strpath), baseline_dir, filename))
 

tests/test_pytest_mpl.py

@@ -39,6 +39,17 @@ def test_succeeds_remote():
     return fig
 
 
+# The following tries an invalid URL first (or at least a URL where the baseline
+# image won't exist), but should succeed with the second mirror.
+@pytest.mark.mpl_image_compare(baseline_dir='http://www.python.org,' + baseline_dir_remote,
+                               filename='test_succeeds_remote.png')
+def test_succeeds_faulty_mirror():
+    fig = plt.figure()
+    ax = fig.add_subplot(1, 1, 1)
+    ax.plot([1, 2, 3])
+    return fig
+
+
 class TestClass(object):
 
     @pytest.mark.mpl_image_compare(baseline_dir=baseline_dir_local)