Merge pull request #39 from dopplershift/output_dir

Control results directory

Author: astrofrog

pytest_mpl/plugin.py

@@ -71,33 +71,47 @@ def pytest_addoption(parser):
     group.addoption('--mpl-baseline-path',
                     help="directory containing baseline images, relative to location where py.test is run", 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')
+    parser.addini('mpl-results-path', help=results_path_help)
+
 
 def pytest_configure(config):
 
     if config.getoption("--mpl") or config.getoption("--mpl-generate-path") is not None:
 
         baseline_dir = config.getoption("--mpl-baseline-path")
         generate_dir = config.getoption("--mpl-generate-path")
+        results_dir = config.getoption("--mpl-results-path") or config.getini("mpl-results-path")
 
-        if baseline_dir is not None and generate_dir is not None:
-            warnings.warn("Ignoring --mpl-baseline-path since --mpl-generate-path is set")
+        if generate_dir is not None:
+            if baseline_dir is not None:
+                warnings.warn("Ignoring --mpl-baseline-path since --mpl-generate-path is set")
+            if results_dir is not None and generate_dir is not None:
+                warnings.warn("Ignoring --mpl-result-path since --mpl-generate-path is set")
 
         if baseline_dir is not None:
             baseline_dir = os.path.abspath(baseline_dir)
         if generate_dir is not None:
             baseline_dir = os.path.abspath(generate_dir)
+        if results_dir is not None:
+            results_dir = os.path.abspath(results_dir)
 
         config.pluginmanager.register(ImageComparison(config,
                                                       baseline_dir=baseline_dir,
-                                                      generate_dir=generate_dir))
+                                                      generate_dir=generate_dir,
+                                                      results_dir=results_dir))
 
 
 class ImageComparison(object):
 
-    def __init__(self, config, baseline_dir=None, generate_dir=None):
+    def __init__(self, config, baseline_dir=None, generate_dir=None, results_dir=None):
         self.config = config
         self.baseline_dir = baseline_dir
         self.generate_dir = generate_dir
+        self.results_dir = results_dir
+        if self.results_dir and not os.path.exists(self.results_dir):
+            os.mkdir(self.results_dir)
 
     def pytest_runtest_setup(self, item):
 
@@ -155,7 +169,7 @@ def item_function_wrapper(*args, **kwargs):
                 if self.generate_dir is None:
 
                     # Save the figure
-                    result_dir = tempfile.mkdtemp()
+                    result_dir = tempfile.mkdtemp(dir=self.results_dir)
                     test_image = os.path.abspath(os.path.join(result_dir, filename))
 
                     fig.savefig(test_image, **savefig_kwargs)
@@ -167,11 +181,9 @@ def item_function_wrapper(*args, **kwargs):
                         baseline_image_ref = os.path.abspath(os.path.join(os.path.dirname(item.fspath.strpath), baseline_dir, filename))
 
                     if not os.path.exists(baseline_image_ref):
-                        raise Exception("""Image file not found for comparison test
-                                        Generated Image:
-                                        \t{test}
-                                        This is expected for new tests.""".format(
-                            test=test_image))
+                        pytest.fail("Image file not found for comparison test. "
+                                    "(This is expected for new tests.)\nGenerated Image: "
+                                    "\n\t{test}".format(test=test_image), pytrace=False)
 
                     # distutils may put the baseline images in non-accessible places,
                     # copy to our tmpdir to be sure to keep them in case of failure
@@ -183,7 +195,7 @@ def item_function_wrapper(*args, **kwargs):
                     if msg is None:
                         shutil.rmtree(result_dir)
                     else:
-                        raise Exception(msg)
+                        pytest.fail(msg, pytrace=False)
 
                 else:
 

tests/test_pytest_mpl.py

@@ -86,6 +86,35 @@ def test_fails(tmpdir):
     assert code == 0
 
 
+TEST_OUTPUT_DIR = """
+import pytest
+import matplotlib.pyplot as plt
+@pytest.mark.mpl_image_compare
+def test_output_dir():
+    fig = plt.figure()
+    ax = fig.add_subplot(1, 1, 1)
+    ax.plot([1, 2, 3])
+    return fig
+"""
+
+
+def test_output_dir(tmpdir):
+    test_file = tmpdir.join('test.py').strpath
+    with open(test_file, 'w') as f:
+        f.write(TEST_OUTPUT_DIR)
+
+    # When we run the test, we should get output images where we specify
+    output_dir = tmpdir.join('test_output_dir').strpath
+    code = subprocess.call('py.test --mpl-results-path={0} --mpl {1}'.format(output_dir, test_file),
+                           shell=True)
+
+    assert code != 0
+    assert os.path.exists(output_dir)
+
+    # Listdir() is to get the random name that the output for the one test is written into
+    assert os.path.exists(os.path.join(output_dir, os.listdir(output_dir)[0], 'test_output_dir.png'))
+
+
 TEST_GENERATE = """
 import pytest
 import matplotlib.pyplot as plt