openobserve/tests/api-testing/support/wait.py

"""Polling and retry helpers for API tests.

`wait_until` replaces `time.sleep(N)`: polls until a predicate is truthy or
the deadline expires.

`retry_on_db_lock` handles transient SQLite 500 responses that appear when
parallel test workers race on write operations against the same DB file.
"""
from __future__ import annotations

import logging
import time
from typing import Any, TypeVar
from collections.abc import Callable

import requests

log = logging.getLogger("o2-api")

T = TypeVar("T")


class WaitTimeout(AssertionError):
    """Raised when `wait_until` exceeds its timeout."""


def wait_until(
    predicate: Callable[[], T],
    *,
    timeout: float = 30.0,
    interval: float = 0.5,
    msg: str = "condition not met",
) -> T:
    """Poll `predicate` until it returns truthy or `timeout` seconds elapse.

    Returns the truthy value the predicate produced (so it can be used as data:
    `result = wait_until(lambda: client.search(...).hits)`).

    Exceptions raised by the predicate are caught and treated as "not yet" until
    the deadline, then re-raised wrapped in `WaitTimeout`. This lets you write
    predicates that do real work (e.g. HTTP calls) without per-call try/except.
    """
    if timeout <= 0:
        raise ValueError("timeout must be > 0")
    if interval <= 0:
        raise ValueError("interval must be > 0")

    deadline = time.monotonic() + timeout
    last_result: Any = None
    last_exc: BaseException | None = None

    while True:
        try:
            last_result = predicate()
            last_exc = None
            if last_result:
                return last_result
        except Exception as e:
            last_exc = e
            last_result = None

        if time.monotonic() >= deadline:
            break
        time.sleep(min(interval, max(0.0, deadline - time.monotonic())))

    if last_exc is not None:
        raise WaitTimeout(f"wait_until: {msg} after {timeout}s; last exception: {last_exc!r}") from last_exc
    raise WaitTimeout(f"wait_until: {msg} after {timeout}s; last value: {last_result!r}")


def retry_on_db_lock(
    fn: Callable[[], requests.Response],
    *,
    max_retries: int = 3,
    base_delay: float = 0.5,
) -> requests.Response:
    """Retry an HTTP call when the server returns a transient SQLite lock error.

    SQLite's single-writer model means parallel xdist workers can race on write
    operations. The server returns 500 with "database is locked" when the
    busy-timeout expires before the lock is released. A short exponential
    backoff is enough for the other writer to commit and release the lock.

    Args:
        fn: Zero-argument callable that performs the HTTP request and returns a
            Response.
        max_retries: Maximum number of additional attempts after the first
            failure (default: 3).
        base_delay: Initial sleep in seconds; doubles on each retry (default:
            0.5s → 1s → 2s).

    Returns:
        The last Response received (success or exhausted retries).

    Example::

        resp = retry_on_db_lock(lambda: session.delete(url))
        assert resp.status_code in (200, 204)
    """
    resp = fn()
    delay = base_delay
    for attempt in range(max_retries):
        if not (resp.status_code == 500 and "database is locked" in resp.text):
            return resp
        log.warning(
            "SQLite database is locked (attempt %d/%d) — retrying in %.1fs",
            attempt + 1,
            max_retries,
            delay,
        )
        time.sleep(delay)
        delay *= 2
        resp = fn()
    return resp