"""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"