comet-ml--opik/sdks/opik_optimizer/tests/unit/utils/test_runtime.py

from __future__ import annotations

import signal
from typing import Any

import pytest

from opik_optimizer import ChatPrompt
from opik_optimizer.base_optimizer import BaseOptimizer
from opik_optimizer.core import runtime
from tests.unit.test_helpers import make_optimization_context


class _ImmediateThread:
    def __init__(
        self, target: Any, args: tuple[Any, ...] = (), daemon: bool | None = None
    ):
        self._target = target
        self._args = args

    def start(self) -> None:
        self._target(*self._args)

    def join(self, timeout: float | None = None) -> None:
        return None


class _NoopTimer:
    def __init__(self, *_args: Any, **_kwargs: Any) -> None:
        pass

    def start(self) -> None:
        return None


class _DummyOptimizer(BaseOptimizer):
    def __init__(self) -> None:
        super().__init__(model="gpt-4o-mini", verbose=0)
        self.calls: list[tuple[Any, str]] = []

    def _finalize_optimization(self, context: Any, status: str = "completed") -> None:
        self.calls.append((context, status))


def test_handle_termination_marks_cancelled(
    monkeypatch: pytest.MonkeyPatch,
) -> None:
    prompt = ChatPrompt(system="sys", user="{question}")
    context = make_optimization_context(prompt)
    optimizer = _DummyOptimizer()

    handlers: dict[int, Any] = {}

    def _fake_getsignal(sig: int) -> Any:
        return f"prev-{sig}"

    def _fake_signal(sig: int, handler: Any) -> None:
        handlers[sig] = handler

    monkeypatch.setattr(signal, "getsignal", _fake_getsignal)
    monkeypatch.setattr(signal, "signal", _fake_signal)
    monkeypatch.setattr(runtime.os, "_exit", lambda _code: None)
    monkeypatch.setattr(runtime.threading, "Thread", _ImmediateThread)
    monkeypatch.setattr(runtime.threading, "Timer", _NoopTimer)

    with runtime.handle_termination(optimizer=optimizer, context=context):
        assert signal.SIGTERM in handlers
        handlers[signal.SIGTERM](signal.SIGTERM, None)

    assert context.should_stop is True
    assert context.finish_reason == "cancelled"
    assert optimizer.calls == [(context, "cancelled")]


def test_candidate_first_aliases_use_history_builder(
    monkeypatch: pytest.MonkeyPatch,
) -> None:
    prompt = ChatPrompt(system="sys", user="{question}")
    context = make_optimization_context(prompt)
    optimizer = _DummyOptimizer()

    calls: list[tuple[str, Any]] = []

    class _HistorySpy:
        def start_round(self, extras: Any | None = None) -> str:
            calls.append(("start_round", extras))
            return "round-handle"

        def record_trial(self, **kwargs: Any) -> None:
            calls.append(("record_trial", kwargs))

        def end_round(self, **kwargs: Any) -> None:
            calls.append(("end_round", kwargs))

    optimizer._history_builder = _HistorySpy()  # type: ignore[assignment]

    round_handle = optimizer.begin_round(context, stage="test")
    candidate_handle = optimizer.start_candidate(
        context, {"candidate": 1}, round_handle=round_handle
    )
    optimizer.finish_candidate(
        context,
        candidate_handle,
        score=0.5,
        round_handle=round_handle,
    )
    optimizer.finish_round(round_handle, context=context, best_score=0.5)

    assert calls[0] == ("start_round", {"stage": "test"})
    assert calls[1][0] == "record_trial"
    assert calls[2][0] == "end_round"