add support for pytest-xdist

Author: ConorMacBride

pytest_mpl/plugin.py

@@ -31,6 +31,7 @@
 import io
 import os
 import json
+import uuid
 import shutil
 import hashlib
 import logging
@@ -216,6 +217,12 @@ def pytest_addoption(parser):
     parser.addini(option, help=msg)
 
 
+class XdistPlugin:
+    def pytest_configure_node(self, node):
+        node.workerinput["pytest_mpl_uid"] = node.config.pytest_mpl_uid
+        node.workerinput["pytest_mpl_results_dir"] = node.config.pytest_mpl_results_dir
+
+
 def pytest_configure(config):
 
     config.addinivalue_line(
@@ -288,12 +295,20 @@ def get_cli_or_ini(name, default=None):
             if not _hash_library_from_cli:
                 hash_library = os.path.abspath(hash_library)
 
+        if not hasattr(config, "workerinput"):
+            uid = uuid.uuid4().hex
+            results_dir_path = results_dir or tempfile.mkdtemp()
+            config.pytest_mpl_uid = uid
+            config.pytest_mpl_results_dir = results_dir_path
+
+        if config.pluginmanager.hasplugin("xdist"):
+            config.pluginmanager.register(XdistPlugin(), name="pytest_mpl_xdist_plugin")
+
         plugin = ImageComparison(
             config,
             baseline_dir=baseline_dir,
             baseline_relative_dir=baseline_relative_dir,
             generate_dir=generate_dir,
-            results_dir=results_dir,
             hash_library=hash_library,
             generate_hash_library=generate_hash_lib,
             generate_summary=generate_summary,
@@ -356,7 +371,6 @@ def __init__(
         baseline_dir=None,
         baseline_relative_dir=None,
         generate_dir=None,
-        results_dir=None,
         hash_library=None,
         generate_hash_library=None,
         generate_summary=None,
@@ -372,7 +386,7 @@ def __init__(
         self.baseline_dir = baseline_dir
         self.baseline_relative_dir = path_is_not_none(baseline_relative_dir)
         self.generate_dir = path_is_not_none(generate_dir)
-        self.results_dir = path_is_not_none(results_dir)
+        self.results_dir = None
         self.hash_library = path_is_not_none(hash_library)
         self._hash_library_from_cli = _hash_library_from_cli  # for backwards compatibility
         self.generate_hash_library = path_is_not_none(generate_hash_library)
@@ -394,11 +408,6 @@ def __init__(
         self.deterministic = deterministic
         self.default_backend = default_backend
 
-        # Generate the containing dir for all test results
-        if not self.results_dir:
-            self.results_dir = Path(tempfile.mkdtemp(dir=self.results_dir))
-        self.results_dir.mkdir(parents=True, exist_ok=True)
-
         # Decide what to call the downloadable results hash library
         if self.hash_library is not None:
             self.results_hash_library_name = self.hash_library.name
@@ -411,6 +420,14 @@ def __init__(
         self._test_stats = None
         self.return_value = {}
 
+    def pytest_sessionstart(self, session):
+        config = session.config
+        if hasattr(config, "workerinput"):
+            config.pytest_mpl_uid = config.workerinput["pytest_mpl_uid"]
+            config.pytest_mpl_results_dir = config.workerinput["pytest_mpl_results_dir"]
+        self.results_dir = Path(config.pytest_mpl_results_dir)
+        self.results_dir.mkdir(parents=True, exist_ok=True)
+
     def get_logger(self):
         # configure a separate logger for this pluggin which is independent
         # of the options that are configured for pytest or for the code that
@@ -932,27 +949,65 @@ def pytest_runtest_call(self, item):  # noqa
                     result._result = None
                     result._excinfo = (type(e), e, e.__traceback__)
 
+    def generate_hash_library_json(self):
+        if hasattr(self.config, "workerinput"):
+            uid = self.config.pytest_mpl_uid
+            worker_id = os.environ.get("PYTEST_XDIST_WORKER")
+            json_file = self.results_dir / f"generated-hashes-xdist-{uid}-{worker_id}.json"
+        else:
+            json_file = Path(self.config.rootdir) / self.generate_hash_library
+            json_file.parent.mkdir(parents=True, exist_ok=True)
+        with open(json_file, 'w') as f:
+            json.dump(self._generated_hash_library, f, indent=2)
+        return json_file
+
     def generate_summary_json(self):
-        json_file = self.results_dir / 'results.json'
+        filename = "results.json"
+        if hasattr(self.config, "workerinput"):
+            uid = self.config.pytest_mpl_uid
+            worker_id = os.environ.get("PYTEST_XDIST_WORKER")
+            filename = f"results-xdist-{uid}-{worker_id}.json"
+        json_file = self.results_dir / filename
         with open(json_file, 'w') as f:
             json.dump(self._test_results, f, indent=2)
         return json_file
 
-    def pytest_unconfigure(self, config):
+    def pytest_sessionfinish(self, session):
         """
         Save out the hash library at the end of the run.
         """
+        config = session.config
+        try:
+            import xdist
+            is_xdist_controller = xdist.is_xdist_controller(session)
+            is_xdist_worker = xdist.is_xdist_worker(session)
+        except ImportError:
+            is_xdist_controller = False
+            is_xdist_worker = False
+        except Exception as e:
+            if "xdist" not in session.config.option:
+                is_xdist_controller = False
+                is_xdist_worker = False
+            else:
+                raise e
+
+        if is_xdist_controller:  # Merge results from workers
+            uid = config.pytest_mpl_uid
+            for worker_hashes in self.results_dir.glob(f"generated-hashes-xdist-{uid}-*.json"):
+                with worker_hashes.open() as f:
+                    self._generated_hash_library.update(json.load(f))
+            for worker_results in self.results_dir.glob(f"results-xdist-{uid}-*.json"):
+                with worker_results.open() as f:
+                    self._test_results.update(json.load(f))
+
         result_hash_library = self.results_dir / (self.results_hash_library_name or "temp.json")
         if self.generate_hash_library is not None:
-            hash_library_path = Path(config.rootdir) / self.generate_hash_library
-            hash_library_path.parent.mkdir(parents=True, exist_ok=True)
-            with open(hash_library_path, "w") as fp:
-                json.dump(self._generated_hash_library, fp, indent=2)
-            if self.results_always:  # Make accessible in results directory
+            hash_library_path = self.generate_hash_library_json()
+            if self.results_always and not is_xdist_worker:  # Make accessible in results directory
                 # Use same name as generated
                 result_hash_library = self.results_dir / hash_library_path.name
                 shutil.copy(hash_library_path, result_hash_library)
-        elif self.results_always and self.results_hash_library_name:
+        elif self.results_always and self.results_hash_library_name and not is_xdist_worker:
             result_hashes = {k: v['result_hash'] for k, v in self._test_results.items()
                              if v['result_hash']}
             if len(result_hashes) > 0:  # At least one hash comparison test
@@ -964,6 +1019,8 @@ def pytest_unconfigure(self, config):
             if 'json' in self.generate_summary:
                 summary = self.generate_summary_json()
                 print(f"A JSON report can be found at: {summary}")
+            if is_xdist_worker:
+                return
             if result_hash_library.exists():  # link to it in the HTML
                 kwargs["hash_library"] = result_hash_library.name
             if 'html' in self.generate_summary:

tests/subtests/helpers.py

@@ -1,4 +1,3 @@
-import os
 import re
 import json
 from pathlib import Path
@@ -8,6 +7,8 @@
 __all__ = ['diff_summary', 'assert_existence', 'patch_summary', 'apply_regex',
            'remove_specific_hashes', 'transform_hashes', 'transform_images']
 
+MIN_EXPECTED_ITEMS = 20  # Rough minimum number of items in a summary to be valid
+
 
 class MatchError(Exception):
     pass
@@ -39,15 +40,26 @@ def diff_summary(baseline, result, baseline_hash_library=None, result_hash_libra
         # Load "correct" baseline hashes
         with open(baseline_hash_library, 'r') as f:
             baseline_hash_library = json.load(f)
+        if len(baseline_hash_library.keys()) < MIN_EXPECTED_ITEMS:
+            raise ValueError(f"baseline_hash_library only has {len(baseline_hash_library.keys())} items")
     else:
         baseline_hash_library = {}
     if result_hash_library and result_hash_library.exists():
         # Load "correct" result hashes
         with open(result_hash_library, 'r') as f:
             result_hash_library = json.load(f)
+        if len(result_hash_library.keys()) < MIN_EXPECTED_ITEMS:
+            raise ValueError(f"result_hash_library only has {len(result_hash_library.keys())} items")
     else:
         result_hash_library = {}
 
+    b = baseline.get("a", baseline)
+    if len(b.keys()) < MIN_EXPECTED_ITEMS:
+        raise ValueError(f"baseline only has {len(b.keys())} items {b}")
+    r = result.get("a", result)
+    if len(r.keys()) < MIN_EXPECTED_ITEMS:
+        raise ValueError(f"result only has {len(r.keys())} items {r}")
+
     # Get test names
     baseline_tests = set(baseline.keys())
     result_tests = set(result.keys())

tests/subtests/test_subtest.py

@@ -48,7 +48,7 @@
 
 
 def run_subtest(baseline_summary_name, tmp_path, args, summaries=None, xfail=True,
-                has_result_hashes=False, generating_hashes=False, testing_hashes=False,
+                has_result_hashes=False, generating_hashes=False, testing_hashes=False, n_xdist_workers=None,
                 update_baseline=UPDATE_BASELINE, update_summary=UPDATE_SUMMARY):
     """ Run pytest (within pytest) and check JSON summary report.
 
@@ -72,6 +72,9 @@ def run_subtest(baseline_summary_name, tmp_path, args, summaries=None, xfail=Tru
         both of `--mpl-hash-library` and `hash_library=` were not.
     testing_hashes : bool, optional, default=False
         Whether the subtest is comparing hashes and therefore needs baseline hashes generated.
+    n_xdist_workers : str or int, optional, default=None
+        Number of xdist workers to use, or "auto" to use all available cores.
+        None will disable xdist. If pytest-xdist is not installed, this will be ignored.
     """
     if update_baseline and update_summary:
         raise ValueError("Cannot enable both `update_baseline` and `update_summary`.")
@@ -109,6 +112,15 @@ def run_subtest(baseline_summary_name, tmp_path, args, summaries=None, xfail=Tru
         shutil.copy(expected_result_hash_library, baseline_hash_library)
         transform_hashes(baseline_hash_library)
 
+    try:
+        import xdist
+        if n_xdist_workers is None:
+            pytest_args += ["-p", "no:xdist"]
+        else:
+            pytest_args += ["-n", str(n_xdist_workers)]
+    except ImportError:
+        pass
+
     # Run the test and record exit status
     status = subprocess.call(pytest_args + mpl_args + args)
 
@@ -206,6 +218,21 @@ def test_html(tmp_path):
     assert (tmp_path / 'results' / 'styles.css').exists()
 
 
+@pytest.mark.parametrize("num_workers", [None, 0, 1, 2])
+def test_html_xdist(request, tmp_path, num_workers):
+    if not request.config.pluginmanager.hasplugin("xdist"):
+        pytest.skip("Skipping: pytest-xdist is not installed")
+    run_subtest('test_results_always', tmp_path,
+                [HASH_LIBRARY_FLAG, BASELINE_IMAGES_FLAG_ABS], summaries=['html'],
+                has_result_hashes=True, n_xdist_workers=num_workers)
+    assert (tmp_path / 'results' / 'fig_comparison.html').exists()
+    assert (tmp_path / 'results' / 'extra.js').exists()
+    assert (tmp_path / 'results' / 'styles.css').exists()
+    if num_workers is not None:
+        assert len(list((tmp_path / 'results').glob('generated-hashes-xdist-*-*.json'))) == 0
+        assert len(list((tmp_path / 'results').glob('results-xdist-*-*.json'))) == num_workers
+
+
 def test_html_hashes_only(tmp_path):
     run_subtest('test_html_hashes_only', tmp_path,
                 [HASH_LIBRARY_FLAG, *HASH_COMPARISON_MODE],
@@ -260,6 +287,24 @@ def test_html_generate(tmp_path):
     assert (tmp_path / 'results' / 'fig_comparison.html').exists()
 
 
+@pytest.mark.parametrize("num_workers", [None, 0, 1, 2])
+def test_html_generate_xdist(request, tmp_path, num_workers):
+    # generating hashes and images; no testing
+    if not request.config.pluginmanager.hasplugin("xdist"):
+        pytest.skip("Skipping: pytest-xdist is not installed")
+    run_subtest('test_html_generate', tmp_path,
+                [rf'--mpl-generate-path={tmp_path}',
+                 rf'--mpl-generate-hash-library={tmp_path / "test_hashes.json"}'],
+                summaries=['html'], xfail=False, has_result_hashes="test_hashes.json",
+                generating_hashes=True, n_xdist_workers=num_workers)
+    assert (tmp_path / 'results' / 'fig_comparison.html').exists()
+    assert (tmp_path / 'results' / 'extra.js').exists()
+    assert (tmp_path / 'results' / 'styles.css').exists()
+    if num_workers is not None:
+        assert len(list((tmp_path / 'results').glob('generated-hashes-xdist-*-*.json'))) == num_workers
+        assert len(list((tmp_path / 'results').glob('results-xdist-*-*.json'))) == num_workers
+
+
 def test_html_generate_images_only(tmp_path):
     # generating images; no testing
     run_subtest('test_html_generate_images_only', tmp_path,

tox.ini

@@ -19,6 +19,7 @@ setenv =
 changedir = .tmp/{envname}
 description = run tests
 deps =
+    pytest-xdist
     mpl20: matplotlib==2.0.*
     mpl21: matplotlib==2.1.*
     mpl22: matplotlib==2.2.*
@@ -58,7 +59,7 @@ commands =
     # Make sure the tests pass with and without --mpl
     # Use -m so pytest skips "subtests" which always apply --mpl
     pytest '{toxinidir}' -m "mpl_image_compare" {posargs}
-    coverage run --source=pytest_mpl -m pytest '{toxinidir}' --mpl
+    coverage run --source=pytest_mpl -m pytest '{toxinidir}' -n auto --mpl
     coverage xml -o '{toxinidir}{/}coverage.xml'
 
 [testenv:codestyle]