"""Utilities for testing with Tkinter"""
import functools


def run_in_tk_mainloop(test_method):
    """Decorator for running a test method with a real Tk mainloop.

    This starts a Tk mainloop before running the test, and stops it
    at the end. This is faster and more robust than the common
    alternative method of calling .update() and/or .update_idletasks().

    Test methods using this must be written as generator functions,
    using "yield" to allow the mainloop to process events and "after"
    callbacks, and then continue the test from that point.

    This also assumes that the test class has a .root attribute,
    which is a tkinter.Tk object.

    For example (from test_sidebar.py):

    @run_test_with_tk_mainloop
    def test_single_empty_input(self):
        self.do_input('\n')
        yield
        self.assert_sidebar_lines_end_with(['>>>', '>>>'])
    """
    @functools.wraps(test_method)
    def new_test_method(self):
        print('start')
        test_generator = test_method(self)
        root = self.root
        # Exceptions raised by self.assert...() need to be raised
        # outside of the after() callback in order for the test
        # harness to capture them.
        exception = None
        firstfail = True
        delay = 1
        def after_callback():
            nonlocal test_generator, exception, firstfail, delay
            try:
                print('next')
                next(test_generator)
            except StopIteration:
                root.quit()
            except Exception as exc:
                if firstfail:
                    print('fail')
                    firstfail = False
                    delay = 8
                    test_generator = test_method(self)
                else:
                    exception = exc
                    root.quit()
            print('delay', delay, flush=True)
            root.after(delay, root.after_idle, after_callback)
        root.after(0, root.after_idle, after_callback)
        root.mainloop()

        if exception:
            raise exception

    return new_test_method