datahub/metadata-ingestion/tests/unit/utilities/test_cooperative_timeout.py

import time

import pytest

from datahub.utilities.cooperative_timeout import (
    CooperativeTimeoutError,
    cooperate,
    cooperative_timeout,
)


def test_cooperate_no_timeout():
    # Called outside of a timeout block, should not do anything.
    cooperate()


def test_cooperate_with_timeout():
    # Set a timeout of 0 seconds, should raise an error immediately
    with pytest.raises(CooperativeTimeoutError), cooperative_timeout(0):
        cooperate()


def test_cooperative_timeout_no_timeout():
    # No timeout set, should not raise an error
    with cooperative_timeout(timeout=None):
        for _ in range(0, 15):
            time.sleep(0.01)
            cooperate()


def test_cooperative_timeout_with_timeout():
    # Set a timeout, and should raise an error after the timeout is hit.
    # It should, however, still run at least one iteration.
    at_least_one_iteration = False
    with pytest.raises(CooperativeTimeoutError), cooperative_timeout(0.5):
        for _ in range(0, 51):
            time.sleep(0.01)
            cooperate()
            at_least_one_iteration = True
    assert at_least_one_iteration