Decorator for running multi threaded unittests

unittest_multi_threaded.py

import unittest
from concurrent.futures import ThreadPoolExecutor


def multi_threaded(*, threads=2):
    def decorator(test_cls):
        for name, test_fn in test_cls.__dict__.copy().items():
            if not (name.startswith("test") and callable(test_fn)):
                continue

            def multi_threaded_test_fn(*args, __test_fn__=test_fn, **kwargs):
                with ThreadPoolExecutor(max_workers=threads) as executor:
                    futures = [
                        executor.submit(__test_fn__, *args, **kwargs)
                        for _ in range(threads)
                    ]
                    futures = iter(futures)
                    try:
                        for future in futures:
                            future.result()
                    finally:
                        for future in futures:
                            future.cancel()

            setattr(test_cls, f"{name}_multi_threaded", multi_threaded_test_fn)

        return test_cls

    return decorator


@multi_threaded(threads=3)
class TestFoo(unittest.TestCase):
    def test_bar(self):
        assert True


if __name__ == "__main__":
    TestFoo.run()