Timeout context manager

timeout.py

import signal
from contextlib import contextmanager
from time import sleep
import pytest


@contextmanager
def timeout(time):
    """Timeout context manager
    
    See: https://www.jujens.eu/posts/en/2018/Jun/02/python-timeout-function/
    """

    def raise_exception(*args):
        raise TimeoutError(f"Timeout of {time} seconds was exceeded")

    signal.signal(signal.SIGALRM, raise_exception)
    signal.alarm(time)

    try:
        yield
    finally:
        signal.signal(signal.SIGALRM, signal.SIG_IGN)


def test_without_timeout():
    with timeout(5):
        for i in range(3):
            sleep(1)


def test_with_timeout():
    with pytest.raises(TimeoutError):
        with timeout(3):
            for i in range(5):
                sleep(1)


def test_with_exception():
    with pytest.raises(RuntimeError):
        with timeout(3):
            sleep(1)
            raise RuntimeError()


if __name__ == "__main__":
    test_without_timeout()
    test_with_timeout()
    test_with_exception()
    print("OK")