put two noise tests in one file

Author: nschloe

test/helpers.py

@@ -10,7 +10,7 @@
 from matplotlib import pyplot as plt
 
 
-def assert_phash(fig, reference_phash):
+def compute_phash(fig):
     # convert to tikz file
     _, tmp_base = tempfile.mkstemp()
     tikz_file = tmp_base + '_tikz.tex'
@@ -82,7 +82,15 @@ def assert_phash(fig, reference_phash):
     png_file = tmp_base + '-1.png'
 
     # compute the phash of the PNG
-    phash = imagehash.phash(Image.open(png_file)).__str__()
+    return (
+        imagehash.phash(Image.open(png_file)).__str__(), png_file, pdf_file,
+        tex_out, ptp_out, mpl_reference, tikz_file
+        )
+
+
+def assert_phash(fig, reference_phash):
+    phash, png_file, pdf_file, tex_out, ptp_out, mpl_reference, tikz_file =\
+        compute_phash(fig)
 
     if reference_phash != phash:
         # Compute the Hamming distance between the two 64-bit numbers

test/test_noise.py

@@ -2,12 +2,13 @@
 #
 import helpers
 
+from numpy.random import randn
+from matplotlib import pyplot as plt
+import numpy as np
+import pytest
 
-def plot():
-    from numpy.random import randn
-    from matplotlib import pyplot as plt
-    import numpy as np
 
+def plot1():
     # Make plot with vertical (default) colorbar
     fig = plt.figure()
     ax = fig.add_subplot(111)
@@ -23,9 +24,39 @@ def plot():
     cbar = fig.colorbar(cax, ticks=[-1, 0, 1])
     # vertically oriented colorbar
     cbar.ax.set_yticklabels(['< -1', '0', '> 1'])
+    return fig
+
+
+def plot2():
+    # Make plot with horizontal colorbar
+    fig = plt.figure()
+    ax = fig.add_subplot(111)
 
+    np.random.seed(123)
+    data = np.clip(np.random.randn(250, 250), -1, 1)
+
+    cax = ax.imshow(data, interpolation='nearest')
+    ax.set_title('Gaussian noise with horizontal colorbar')
+
+    cbar = fig.colorbar(cax, ticks=[-1, 0, 1], orientation='horizontal')
+    # horizontal colorbar
+    cbar.ax.set_xticklabels(['Low', 'Medium', 'High'])
     return fig
 
 
-def test():
-    helpers.assert_phash(plot(), 'f55a0bb503fd0354')
+@pytest.mark.parametrize(
+    'plot, phash', [
+        (plot1, 'fda50b5a0b5a835a'),
+        (plot2, '7f56a913a9816ba1'),
+    ]
+    )
+def test(plot, phash):
+    helpers.assert_phash(plot(), phash)
+    return
+
+
+if __name__ == '__main__':
+    # plot1()
+    # plt.show()
+    phash, _, _, _, _, _, _ = helpers.compute_phash(plot2())
+    print(phash)