Support yield to save a figure multiple times

I have some tests for figures which use widgets to update the figure. It would be really helpful if I could yield from the test with a figure and then update it and yield from it again or return the figure the final time.

(I am not 100% sure if that's possible with pytest)

I'm quite sure pytest doesn't support yielding inside a test function, however, this workaround seems the easiest solution. It dynamically creates test methods in the test class. You just give it a figure generator object and it the number of times it yields. We could add it to the sphinx website as an example?

import pytest
import matplotlib.pyplot as plt


def yielder(test_function, n_tests, **kwargs):

    def run_test():
        @pytest.mark.mpl_image_compare(**kwargs)
        def test(*args):
            return next(test_function)
        return test

    class MetaYielder(type):
        def __init__(cls, name, bases, dct):
            for i in range(1, n_tests+1):
                setattr(cls, f"test_{i}", run_test())
    return MetaYielder


def yielding_test():
    fig = plt.figure()
    ax = fig.add_subplot(1, 1, 1)
    ax.plot([1, 2, 1, 1])
    yield fig
    ax.plot([1, 1, 2, 1])
    yield fig


class TestYield(metaclass=yielder(yielding_test(), 2)):
    pass