"""Tests for OutputManager (auto-save workflow outputs)"""
import pytest
from pytest_httpx import HTTPXMock
from src.auth.base import NoAuth
from src.orchestrators.output import OutputManager
class TestOutputManager:
"""Test OutputManager functionality"""
def test_init(self, tmp_path):
"""Test OutputManager initialization"""
auth = NoAuth("http://localhost:8188")
output_root = tmp_path / "executions" / "generations"
manager = OutputManager(auth=auth, output_root=output_root)
assert manager.output_root == output_root
assert output_root.exists() # Should create directory
def test_get_workflow_output_dir(self, tmp_path):
"""Test get_workflow_output_dir creates subdirectory"""
auth = NoAuth("http://localhost:8188")
output_root = tmp_path / "executions" / "generations"
manager = OutputManager(auth=auth, output_root=output_root)
workflow_dir = manager.get_workflow_output_dir("generate_image")
assert workflow_dir == output_root / "generate_image"
def test_list_workflow_outputs_empty(self, tmp_path):
"""Test list_workflow_outputs with no files"""
auth = NoAuth("http://localhost:8188")
output_root = tmp_path / "executions" / "generations"
manager = OutputManager(auth=auth, output_root=output_root)
outputs = manager.list_workflow_outputs("generate_image")
assert outputs == []
@pytest.mark.asyncio
async def test_save_workflow_output_basic(self, tmp_path, httpx_mock: HTTPXMock):
"""Test save_workflow_output with basic image"""
auth = NoAuth("http://localhost:8188")
output_root = tmp_path / "executions" / "generations"
# Mock ComfyUI asset fetch
fake_image_data = b"fake-png-data"
httpx_mock.add_response(
url="http://localhost:8188/view?filename=test.png&type=output",
content=fake_image_data,
headers={"content-type": "image/png"},
)
manager = OutputManager(auth=auth, output_root=output_root)
asset_info = {
"filename": "test.png",
"subfolder": "",
"folder_type": "output",
}
success, saved_path, error = await manager.save_workflow_output(
asset_info=asset_info,
workflow_id="generate_image",
web_optimize=False, # Disable to avoid PIL dependency in test
)
assert success is True
assert error is None
assert saved_path is not None
assert saved_path.exists()
assert saved_path.parent.name == "generate_image"
assert "generate_image" in saved_path.name
assert saved_path.suffix == ".png"
# Verify content
with open(saved_path, "rb") as f:
assert f.read() == fake_image_data
@pytest.mark.asyncio
async def test_save_workflow_output_missing_filename(self, tmp_path):
"""Test save_workflow_output with missing filename"""
auth = NoAuth("http://localhost:8188")
output_root = tmp_path / "executions" / "generations"
manager = OutputManager(auth=auth, output_root=output_root)
asset_info = {
"subfolder": "",
"folder_type": "output",
}
success, saved_path, error = await manager.save_workflow_output(
asset_info=asset_info,
workflow_id="generate_image",
web_optimize=False,
)
assert success is False
assert saved_path is None
assert "No filename" in error
@pytest.mark.asyncio
async def test_save_workflow_output_fetch_failure(self, tmp_path, httpx_mock: HTTPXMock):
"""Test save_workflow_output when ComfyUI fetch fails"""
auth = NoAuth("http://localhost:8188")
output_root = tmp_path / "executions" / "generations"
# Mock 404 response
httpx_mock.add_response(
url="http://localhost:8188/view?filename=missing.png&type=output",
status_code=404,
)
manager = OutputManager(auth=auth, output_root=output_root)
asset_info = {
"filename": "missing.png",
"subfolder": "",
"folder_type": "output",
}
success, saved_path, error = await manager.save_workflow_output(
asset_info=asset_info,
workflow_id="generate_image",
web_optimize=False,
)
assert success is False
assert saved_path is None
assert "Failed to get asset" in error or "HTTP 404" in error
@pytest.mark.asyncio
async def test_save_workflow_output_with_subfolder(self, tmp_path, httpx_mock: HTTPXMock):
"""Test save_workflow_output with subfolder"""
auth = NoAuth("http://localhost:8188")
output_root = tmp_path / "executions" / "generations"
fake_image_data = b"fake-image-with-subfolder"
httpx_mock.add_response(
url="http://localhost:8188/view?filename=output.png&subfolder=my_folder&type=output",
content=fake_image_data,
headers={"content-type": "image/png"},
)
manager = OutputManager(auth=auth, output_root=output_root)
asset_info = {
"filename": "output.png",
"subfolder": "my_folder",
"folder_type": "output",
}
success, saved_path, error = await manager.save_workflow_output(
asset_info=asset_info,
workflow_id="test_workflow",
web_optimize=False,
)
assert success is True
assert error is None
assert saved_path.exists()
with open(saved_path, "rb") as f:
assert f.read() == fake_image_data
def test_list_workflow_outputs_with_files(self, tmp_path):
"""Test list_workflow_outputs returns sorted files"""
auth = NoAuth("http://localhost:8188")
output_root = tmp_path / "executions" / "generations"
workflow_dir = output_root / "generate_image"
workflow_dir.mkdir(parents=True)
# Create test files with different timestamps
file1 = workflow_dir / "image-001.png"
file2 = workflow_dir / "image-002.png"
file3 = workflow_dir / "image-003.png"
file1.write_text("test1")
import time
time.sleep(0.01)
file2.write_text("test2")
time.sleep(0.01)
file3.write_text("test3")
manager = OutputManager(auth=auth, output_root=output_root)
outputs = manager.list_workflow_outputs("generate_image")
# Should be sorted by most recent first
assert len(outputs) == 3
assert outputs[0].name == "image-003.png" # Most recent
assert outputs[-1].name == "image-001.png" # Oldest
