tests/test_executor.py

"""ZeroGPU command serialization and parent-side error restoration."""

from __future__ import annotations

import pickle
from pathlib import Path

import pytest

from core.errors import GatewayError
from core.executor import (
    CommandError,
    CommandResult,
    InferenceCommand,
    execute_inference,
)


def test_inference_command_is_pickle_safe(tmp_path: Path) -> None:
    command = InferenceCommand(
        model_name="flux",
        method_name="generate",
        arguments={"prompt": "city", "output_path": tmp_path / "image.png"},
        request_id="request-1",
        duration_seconds=180,
    )

    restored = pickle.loads(pickle.dumps(command))

    assert restored == command


def test_serialized_worker_error_restores_gateway_error(monkeypatch) -> None:
    command = InferenceCommand(
        model_name="flux",
        method_name="generate",
        arguments={},
        request_id="request-1",
        duration_seconds=180,
    )
    monkeypatch.setattr(
        "core.executor._invoke_large",
        lambda _: CommandResult(
            error=CommandError("flux ran out of memory", 507, "out_of_memory")
        ),
    )

    with pytest.raises(GatewayError) as captured:
        execute_inference(command)

    assert captured.value.status_code == 507
    assert captured.value.code == "out_of_memory"