Skip to content

connection_utils

drunc.grpc_testing_tools.connection_utils

Functions

wait_for(condition, expected_value, timeout=10.0, poll_interval=0.1)

Wait for a condition to return an expected value within a timeout period.

Repeatedly evaluates the condition callable until it returns the expected value or the timeout is reached. Useful for polling operations in tests and async workflows.

Parameters:

Name Type Description Default
condition Callable[[], Any]

Callable that returns a value to check. Should take no arguments.

 required
expected_value Any

The value to wait for. If None, waits for any truthy value. Can also be a tuple of acceptable values.

 required
timeout float

Maximum time to wait in seconds before raising TimeoutError.

 10.0
poll_interval float

Time in seconds between condition evaluations.

 0.1

Returns:

Type Description
Optional[Any]

The value returned by condition when it matches expected_value.

Raises:

Type Description
TimeoutError

If the condition doesn't return expected_value within timeout.
Source code in drunc/grpc_testing_tools/connection_utils.py 
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
def wait_for(
    condition: Callable[[], Any],
    expected_value: Any,
    timeout: float = 10.0,
    poll_interval: float = 0.1,
) -> Optional[Any]:
    """
    Wait for a condition to return an expected value within a timeout period.

    Repeatedly evaluates the condition callable until it returns the expected value
    or the timeout is reached. Useful for polling operations in tests and async workflows.

    Args:
        condition: Callable that returns a value to check. Should take no arguments.
        expected_value: The value to wait for. If None, waits for any truthy value.
                       Can also be a tuple of acceptable values.
        timeout: Maximum time to wait in seconds before raising TimeoutError.
        poll_interval: Time in seconds between condition evaluations.

    Returns:
        The value returned by condition when it matches expected_value.

    Raises:
        TimeoutError: If the condition doesn't return expected_value within timeout.
    """
    start_time = time.time()
    last_value = None

    while time.time() - start_time < timeout:
        try:
            last_value = condition()
        except Exception:
            time.sleep(poll_interval)
            continue

        # Handle callable expected_value (predicate function)
        if callable(expected_value):
            if expected_value(last_value):
                return last_value
        # Handle tuple of acceptable values
        elif isinstance(expected_value, tuple):
            if last_value in expected_value:
                return last_value
        # Handle None as "wait for any truthy value"
        elif expected_value is None:
            if last_value:
                return last_value
        # Handle direct value comparison
        else:
            if last_value == expected_value:
                return last_value

        time.sleep(poll_interval)

    return None