from unittest import mock


class CatchSignal:
    """
    Context manager class for testing that Django signals are called.
    https://hakibenita.com/how-to-test-django-signals-like-a-pro

    with CatchSignal(signal_name) as handler:
        print('do things that send signals')

        handler.assert_called_once()
        signal_args = handler.call_args[1]
        self.assertEqual(signal_args['sender'], Model)
        self.assertTrue(signal_args['created'])
    """

    def __init__(self, signal):
        self.signal = signal
        self.handler = mock.Mock()

    def __enter__(self):
        self.signal.connect(self.handler)
        return self.handler

    def __exit__(self, exc_type, exc_value, tb):
        self.signal.disconnect(self.handler)
