fix(ml): comfyui_wait_result no sale prematuro en jobs de video/3D
Exige outputs no vacios (no solo status terminal) para dar por completado un prompt: en jobs pesados ComfyUI marca la entry de /history como terminada antes de poblar outputs, lo que devolvia un dict vacio mientras el job seguia en GPU. Ahora sigue sondeando hasta que los outputs aparecen o hasta agotar el timeout. Timeout default 180s -> 600s (cubre video/3D) y timeout HTTP por-request acotado a 30s. Firma y contrato de retorno intactos. Tests nuevos (mock urllib CI-safe + live opcional contra /history real): golden, regresion del bug, edge imagen corta, timeout y error. v1.0.0 -> 1.1.0. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
This commit is contained in:
@@ -0,0 +1,179 @@
|
||||
"""Tests para comfyui_wait_result (funcion impura: polling de /history).
|
||||
|
||||
La mayoria mockea urllib.request.urlopen para correr sin GPU ni red (CI-safe).
|
||||
Un test "live" extra valida contra el /history real de un job de video ya
|
||||
completado, y se salta limpiamente si el servidor ComfyUI no responde o el job
|
||||
ya no esta en el history (p. ej. servidor reiniciado).
|
||||
"""
|
||||
|
||||
import json
|
||||
import os
|
||||
import sys
|
||||
from unittest import mock
|
||||
|
||||
import pytest
|
||||
|
||||
sys.path.insert(0, os.path.dirname(__file__))
|
||||
sys.path.insert(0, os.path.join(os.path.dirname(__file__), "..", ".."))
|
||||
|
||||
from ml.comfyui_wait_result import comfyui_wait_result
|
||||
|
||||
PID = "test-pid-1"
|
||||
|
||||
|
||||
class _FakeResp:
|
||||
"""Context manager minimo que imita la respuesta de urllib.request.urlopen."""
|
||||
|
||||
def __init__(self, payload):
|
||||
self._payload = json.dumps(payload).encode()
|
||||
|
||||
def read(self):
|
||||
return self._payload
|
||||
|
||||
def __enter__(self):
|
||||
return self
|
||||
|
||||
def __exit__(self, *_):
|
||||
return False
|
||||
|
||||
|
||||
def _seq_urlopen(payloads):
|
||||
"""Fake urlopen que devuelve cada payload en orden; repite el ultimo."""
|
||||
state = {"i": 0}
|
||||
|
||||
def _open(url, timeout=None):
|
||||
i = min(state["i"], len(payloads) - 1)
|
||||
state["i"] += 1
|
||||
return _FakeResp(payloads[i])
|
||||
|
||||
return _open
|
||||
|
||||
|
||||
def _entry(status_str="success", completed=True, outputs=None):
|
||||
return {
|
||||
PID: {
|
||||
"status": {"status_str": status_str, "completed": completed, "messages": []},
|
||||
"outputs": outputs if outputs is not None else {},
|
||||
}
|
||||
}
|
||||
|
||||
|
||||
# --- Golden: espera a outputs no vacios en un job largo (entry tarda en aparecer) ---
|
||||
def test_golden_espera_a_que_aparezcan_outputs():
|
||||
payloads = [
|
||||
{}, # aun en cola: entry ausente
|
||||
{}, # sigue en cola
|
||||
_entry(
|
||||
outputs={
|
||||
"79": {
|
||||
"images": [
|
||||
{"filename": "vid.mp4", "subfolder": "", "type": "output"}
|
||||
],
|
||||
"animated": [True],
|
||||
}
|
||||
}
|
||||
),
|
||||
]
|
||||
with mock.patch(
|
||||
"ml.comfyui_wait_result.urllib.request.urlopen", _seq_urlopen(payloads)
|
||||
), mock.patch("ml.comfyui_wait_result.time.sleep", lambda *_: None):
|
||||
out = comfyui_wait_result(PID, timeout=5, poll_interval=0.001)
|
||||
assert out, "debe devolver outputs no vacios"
|
||||
assert out["79"]["images"][0]["filename"] == "vid.mp4"
|
||||
|
||||
|
||||
# --- Regresion del bug: 'done' con outputs vacios NO debe salir prematuro ---
|
||||
def test_no_sale_prematuro_cuando_done_pero_outputs_vacios():
|
||||
payloads = [
|
||||
_entry(outputs={}), # estado terminal pero outputs aun vacios (caso del bug)
|
||||
_entry(outputs={}), # sigue vacio mientras la GPU trabaja
|
||||
_entry(outputs={"9": {"images": [{"filename": "img.png"}]}}), # ya poblado
|
||||
]
|
||||
with mock.patch(
|
||||
"ml.comfyui_wait_result.urllib.request.urlopen", _seq_urlopen(payloads)
|
||||
), mock.patch("ml.comfyui_wait_result.time.sleep", lambda *_: None):
|
||||
out = comfyui_wait_result(PID, timeout=5, poll_interval=0.001)
|
||||
assert out == {"9": {"images": [{"filename": "img.png"}]}}
|
||||
|
||||
|
||||
# --- Edge: imagen corta ya completada en el primer sondeo -> retorno inmediato ---
|
||||
def test_imagen_corta_devuelve_en_primer_poll():
|
||||
payloads = [_entry(outputs={"9": {"images": [{"filename": "apple.png"}]}})]
|
||||
opener = _seq_urlopen(payloads)
|
||||
with mock.patch(
|
||||
"ml.comfyui_wait_result.urllib.request.urlopen", opener
|
||||
), mock.patch("ml.comfyui_wait_result.time.sleep", lambda *_: None):
|
||||
out = comfyui_wait_result(PID, timeout=5, poll_interval=0.001)
|
||||
assert out["9"]["images"][0]["filename"] == "apple.png"
|
||||
|
||||
|
||||
# --- Edge: estado 'completed' sin status_str success tambien vale (si hay outputs) ---
|
||||
def test_completed_sin_status_str_success():
|
||||
payloads = [_entry(status_str=None, completed=True, outputs={"9": {"images": [{"filename": "x.png"}]}})]
|
||||
with mock.patch(
|
||||
"ml.comfyui_wait_result.urllib.request.urlopen", _seq_urlopen(payloads)
|
||||
), mock.patch("ml.comfyui_wait_result.time.sleep", lambda *_: None):
|
||||
out = comfyui_wait_result(PID, timeout=5, poll_interval=0.001)
|
||||
assert out["9"]["images"][0]["filename"] == "x.png"
|
||||
|
||||
|
||||
# --- Error: timeout genuino (prompt nunca aparece) -> TimeoutError, no cuelga ---
|
||||
def test_timeout_prompt_inexistente_lanza_timeout():
|
||||
payloads = [{}] # entry siempre ausente
|
||||
with mock.patch(
|
||||
"ml.comfyui_wait_result.urllib.request.urlopen", _seq_urlopen(payloads)
|
||||
):
|
||||
with pytest.raises(TimeoutError):
|
||||
comfyui_wait_result(PID, timeout=0.05, poll_interval=0.02)
|
||||
|
||||
|
||||
# --- Error: ejecucion fallida (status_str 'error') -> RuntimeError ---
|
||||
def test_status_error_lanza_runtimeerror():
|
||||
payloads = [
|
||||
{
|
||||
PID: {
|
||||
"status": {
|
||||
"status_str": "error",
|
||||
"completed": False,
|
||||
"messages": [["execution_error", {"exception_message": "OOM"}]],
|
||||
},
|
||||
"outputs": {},
|
||||
}
|
||||
}
|
||||
]
|
||||
with mock.patch(
|
||||
"ml.comfyui_wait_result.urllib.request.urlopen", _seq_urlopen(payloads)
|
||||
), mock.patch("ml.comfyui_wait_result.time.sleep", lambda *_: None):
|
||||
with pytest.raises(RuntimeError):
|
||||
comfyui_wait_result(PID, timeout=5, poll_interval=0.001)
|
||||
|
||||
|
||||
# --- Live (opcional): job de video ya completado en el /history real ---
|
||||
LIVE_SERVER = "127.0.0.1:8188"
|
||||
LIVE_PID = "8a278988-8a94-4225-add3-88a406f7101c" # LTX T2V del report 0105
|
||||
|
||||
|
||||
def _server_up(server=LIVE_SERVER):
|
||||
import urllib.request
|
||||
|
||||
try:
|
||||
urllib.request.urlopen(f"http://{server}/history/__probe__", timeout=2).read()
|
||||
return True
|
||||
except Exception:
|
||||
return False
|
||||
|
||||
|
||||
@pytest.mark.skipif(not _server_up(), reason="ComfyUI no responde en 127.0.0.1:8188")
|
||||
def test_live_job_video_completado_devuelve_outputs():
|
||||
import urllib.request
|
||||
|
||||
h = json.loads(
|
||||
urllib.request.urlopen(
|
||||
f"http://{LIVE_SERVER}/history/{LIVE_PID}", timeout=3
|
||||
).read()
|
||||
)
|
||||
if LIVE_PID not in h:
|
||||
pytest.skip("job 8a278988 ya no esta en /history (servidor reiniciado)")
|
||||
out = comfyui_wait_result(LIVE_PID, server=LIVE_SERVER, timeout=10)
|
||||
assert out, "un job ya completado debe devolver outputs no vacios y rapido"
|
||||
assert any(v.get("images") for v in out.values()), "debe traer la clave 'images'"
|
||||
Reference in New Issue
Block a user