TST: Parametrize tests that were using yield.

matplotlib · NelleV · Oct 22, 2016 · Oct 20, 2016 · Oct 20, 2016 · Oct 21, 2016
commit ba30f11c0e1bc84dff2597ecf97239442d61ac79
diff --git a/lib/matplotlib/tests/test_animation.py b/lib/matplotlib/tests/test_animation.py
@@ -6,7 +6,10 @@
 import os
 import sys
 import tempfile
+
 import numpy as np
+import pytest
+
 import matplotlib as mpl
 from matplotlib import pyplot as plt
 from matplotlib import animation
@@ -91,23 +94,25 @@ def isAvailable(self):
         return True
 
 
-WRITER_OUTPUT = dict(ffmpeg='mp4', ffmpeg_file='mp4',
-                     mencoder='mp4', mencoder_file='mp4',
-                     avconv='mp4', avconv_file='mp4',
-                     imagemagick='gif', imagemagick_file='gif',
-                     null='null')
+WRITER_OUTPUT = [
+    ('ffmpeg', 'mp4'),
+    ('ffmpeg_file', 'mp4'),
+    ('mencoder', 'mp4'),
+    ('mencoder_file', 'mp4'),
+    ('avconv', 'mp4'),
+    ('avconv_file', 'mp4'),
+    ('imagemagick', 'gif'),
+    ('imagemagick_file', 'gif'),
+    ('null', 'null')
+]
 
 
 # Smoke test for saving animations.  In the future, we should probably
 # design more sophisticated tests which compare resulting frames a-la
 # matplotlib.testing.image_comparison
-def test_save_animation_smoketest():
-    for writer, extension in six.iteritems(WRITER_OUTPUT):
-        yield check_save_animation, writer, extension
-
-
 @cleanup
-def check_save_animation(writer, extension='mp4'):
+@pytest.mark.parametrize('writer, extension', WRITER_OUTPUT)
+def test_save_animation_smoketest(writer, extension):
     try:
         # for ImageMagick the rcparams must be patched to account for
         # 'convert' being a built in MS tool, not the imagemagick

diff --git a/lib/matplotlib/tests/test_axes.py b/lib/matplotlib/tests/test_axes.py
@@ -18,6 +18,7 @@
 from numpy import ma
 from numpy import arange
 from cycler import cycler
+import pytest
 
 import warnings
 
@@ -4269,13 +4270,7 @@ def test_violin_point_mass():
     plt.violinplot(np.array([0, 0]))
 
 
-def _eb_succes_helper(ax, x, y, xerr=None, yerr=None):
-    eb = ax.errorbar(x, y, xerr=xerr, yerr=yerr)
-    eb.remove()
-
-
-@cleanup
-def test_errorbar_inputs_shotgun():
+def generate_errorbar_inputs():
     base_xy = cycler('x', [np.arange(5)]) + cycler('y', [np.ones((5, ))])
     err_cycler = cycler('err', [1,
                                 [1, 1, 1, 1, 1],
@@ -4296,15 +4291,17 @@ def test_errorbar_inputs_shotgun():
     yerr_only = base_xy * yerr_cy
     both_err = base_xy * yerr_cy * xerr_cy
 
-    test_cyclers = xerr_only, yerr_only, both_err, empty
+    test_cyclers = chain(xerr_only, yerr_only, both_err, empty)
+
+    return test_cyclers
+
 
+@cleanup
+@pytest.mark.parametrize('kwargs', generate_errorbar_inputs())
+def test_errorbar_inputs_shotgun(kwargs):
     ax = plt.gca()
-    # should do this as a generative test, but @cleanup seems to break that
-    # for p in chain(*test_cyclers):
-    #     yield (_eb_succes_helper, ax) + tuple(p.get(k, None) for
-    #                                          k in ['x', 'y', 'xerr', 'yerr'])
-    for p in chain(*test_cyclers):
-        _eb_succes_helper(ax, **p)
+    eb = ax.errorbar(**kwargs)
+    eb.remove()
 
 
 @cleanup
@@ -4386,8 +4383,8 @@ def test_axes_margins():
     assert ax.get_ybound() == (-0.5, 9.5)
 
 
-@cleanup
-def test_remove_shared_axes():
+@pytest.fixture(params=['x', 'y'])
+def shared_axis_remover(request):
     def _helper_x(ax):
         ax2 = ax.twinx()
         ax2.remove()
@@ -4402,26 +4399,43 @@ def _helper_y(ax):
         r = ax.yaxis.get_major_locator()()
         assert r[-1] > 14
 
+    if request.param == 'x':
+        return _helper_x
+    elif request.param == 'y':
+        return _helper_y
+    else:
+        assert False, 'Request param %s is invalid.' % (request.param, )
+
+
+@pytest.fixture(params=['gca', 'subplots', 'subplots_shared', 'add_axes'])
+def shared_axes_generator(request):
     # test all of the ways to get fig/ax sets
-    fig = plt.figure()
-    ax = fig.gca()
-    yield _helper_x, ax
-    yield _helper_y, ax
+    if request.param == 'gca':
+        fig = plt.figure()
+        ax = fig.gca()
+    elif request.param == 'subplots':
+        fig, ax = plt.subplots()
+    elif request.param == 'subplots_shared':
+        fig, ax_lst = plt.subplots(2, 2, sharex='all', sharey='all')
+        ax = ax_lst[0][0]
+    elif request.param == 'add_axes':
+        fig = plt.figure()
+        ax = fig.add_axes([.1, .1, .8, .8])
+    else:
+        assert False, 'Request param %s is invalid.' % (request.param, )
 
-    fig, ax = plt.subplots()
-    yield _helper_x, ax
-    yield _helper_y, ax
+    return fig, ax
 
-    fig, ax_lst = plt.subplots(2, 2, sharex='all', sharey='all')
-    ax = ax_lst[0][0]
-    yield _helper_x, ax
-    yield _helper_y, ax
 
-    fig = plt.figure()
-    ax = fig.add_axes([.1, .1, .8, .8])
-    yield _helper_x, ax
-    yield _helper_y, ax
+@cleanup
+def test_remove_shared_axes(shared_axes_generator, shared_axis_remover):
+    # test all of the ways to get fig/ax sets
+    fig, ax = shared_axes_generator
+    shared_axis_remover(ax)
+
 
+@cleanup
+def test_remove_shared_axes_relim():
     fig, ax_lst = plt.subplots(2, 2, sharex='all', sharey='all')
     ax = ax_lst[0][0]
     orig_xlim = ax_lst[0][1].get_xlim()

diff --git a/lib/matplotlib/tests/test_cbook.py b/lib/matplotlib/tests/test_cbook.py
@@ -12,6 +12,7 @@
 from numpy.testing.utils import (assert_array_equal, assert_approx_equal,
                                  assert_array_almost_equal)
 from nose.tools import raises, assert_raises
+import pytest
 
 import matplotlib.cbook as cbook
 import matplotlib.colors as mcolors
@@ -343,65 +344,61 @@ def test_sanitize_sequence():
     assert k == cbook.sanitize_sequence(k)
 
 
-def _kwarg_norm_helper(inp, expected, kwargs_to_norm, warn_count=0):
-
+fail_mapping = (
+    ({'a': 1}, {'forbidden': ('a')}),
+    ({'a': 1}, {'required': ('b')}),
+    ({'a': 1, 'b': 2}, {'required': ('a'), 'allowed': ()})
+)
+
+warn_passing_mapping = (
+    ({'a': 1, 'b': 2}, {'a': 1}, {'alias_mapping': {'a': ['b']}}, 1),
+    ({'a': 1, 'b': 2}, {'a': 1},
+     {'alias_mapping': {'a': ['b']}, 'allowed': ('a',)}, 1),
+    ({'a': 1, 'b': 2}, {'a': 2}, {'alias_mapping': {'a': ['a', 'b']}}, 1),
+    ({'a': 1, 'b': 2, 'c': 3}, {'a': 1, 'c': 3},
+     {'alias_mapping': {'a': ['b']}, 'required': ('a', )}, 1),
+)
+
+pass_mapping = (
+    ({'a': 1, 'b': 2}, {'a': 1, 'b': 2}, {}),
+    ({'b': 2}, {'a': 2}, {'alias_mapping': {'a': ['a', 'b']}}),
+    ({'b': 2}, {'a': 2},
+     {'alias_mapping': {'a': ['b']}, 'forbidden': ('b', )}),
+    ({'a': 1, 'c': 3}, {'a': 1, 'c': 3},
+     {'required': ('a', ), 'allowed': ('c', )}),
+    ({'a': 1, 'c': 3}, {'a': 1, 'c': 3},
+     {'required': ('a', 'c'), 'allowed': ('c', )}),
+    ({'a': 1, 'c': 3}, {'a': 1, 'c': 3},
+     {'required': ('a', 'c'), 'allowed': ('a', 'c')}),
+    ({'a': 1, 'c': 3}, {'a': 1, 'c': 3},
+     {'required': ('a', 'c'), 'allowed': ()}),
+    ({'a': 1, 'c': 3}, {'a': 1, 'c': 3}, {'required': ('a', 'c')}),
+    ({'a': 1, 'c': 3}, {'a': 1, 'c': 3}, {'allowed': ('a', 'c')}),
+)
+
+
+@pytest.mark.parametrize('inp, kwargs_to_norm', fail_mapping)
+def test_normalize_kwargs_fail(inp, kwargs_to_norm):
+    with pytest.raises(TypeError):
+        cbook.normalize_kwargs(inp, **kwargs_to_norm)
+
+
+@pytest.mark.parametrize('inp, expected, kwargs_to_norm, warn_count',
+                         warn_passing_mapping)
+def test_normalize_kwargs_warn(inp, expected, kwargs_to_norm, warn_count):
     with warnings.catch_warnings(record=True) as w:
         warnings.simplefilter("always")
         assert expected == cbook.normalize_kwargs(inp, **kwargs_to_norm)
         assert len(w) == warn_count
 
 
-def _kwarg_norm_fail_helper(inp, kwargs_to_norm):
-    assert_raises(TypeError, cbook.normalize_kwargs, inp, **kwargs_to_norm)
-
-
-def test_normalize_kwargs():
-    fail_mapping = (
-        ({'a': 1}, {'forbidden': ('a')}),
-        ({'a': 1}, {'required': ('b')}),
-        ({'a': 1, 'b': 2}, {'required': ('a'), 'allowed': ()})
-    )
-
-    for inp, kwargs in fail_mapping:
-        yield _kwarg_norm_fail_helper, inp, kwargs
-
-    warn_passing_mapping = (
-        ({'a': 1, 'b': 2}, {'a': 1}, {'alias_mapping': {'a': ['b']}}, 1),
-        ({'a': 1, 'b': 2}, {'a': 1}, {'alias_mapping': {'a': ['b']},
-                                      'allowed': ('a',)}, 1),
-        ({'a': 1, 'b': 2}, {'a': 2}, {'alias_mapping': {'a': ['a', 'b']}}, 1),
-
-        ({'a': 1, 'b': 2, 'c': 3}, {'a': 1, 'c': 3},
-         {'alias_mapping': {'a': ['b']}, 'required': ('a', )}, 1),
-
-    )
-
-    for inp, exp, kwargs, wc in warn_passing_mapping:
-        yield _kwarg_norm_helper, inp, exp, kwargs, wc
-
-    pass_mapping = (
-        ({'a': 1, 'b': 2}, {'a': 1, 'b': 2}, {}),
-        ({'b': 2}, {'a': 2}, {'alias_mapping': {'a': ['a', 'b']}}),
-        ({'b': 2}, {'a': 2}, {'alias_mapping': {'a': ['b']},
-                              'forbidden': ('b', )}),
-
-        ({'a': 1, 'c': 3}, {'a': 1, 'c': 3}, {'required': ('a', ),
-                                              'allowed': ('c', )}),
-
-        ({'a': 1, 'c': 3}, {'a': 1, 'c': 3}, {'required': ('a', 'c'),
-                                              'allowed': ('c', )}),
-        ({'a': 1, 'c': 3}, {'a': 1, 'c': 3}, {'required': ('a', 'c'),
-                                              'allowed': ('a', 'c')}),
-        ({'a': 1, 'c': 3}, {'a': 1, 'c': 3}, {'required': ('a', 'c'),
-                                              'allowed': ()}),
-
-        ({'a': 1, 'c': 3}, {'a': 1, 'c': 3}, {'required': ('a', 'c')}),
-        ({'a': 1, 'c': 3}, {'a': 1, 'c': 3}, {'allowed': ('a', 'c')}),
-
-    )
-
-    for inp, exp, kwargs in pass_mapping:
-        yield _kwarg_norm_helper, inp, exp, kwargs
+@pytest.mark.parametrize('inp, expected, kwargs_to_norm',
+                         pass_mapping)
+def test_normalize_kwargs_pass(inp, expected, kwargs_to_norm):
+    with warnings.catch_warnings(record=True) as w:
+        warnings.simplefilter("always")
+        assert expected == cbook.normalize_kwargs(inp, **kwargs_to_norm)
+        assert len(w) == 0
 
 
 def test_to_prestep():