110 lines
4.5 KiB
Python
110 lines
4.5 KiB
Python
import asyncio
|
|
import sys
|
|
import os
|
|
from unittest.mock import MagicMock, patch
|
|
|
|
sys.path.append(os.getcwd())
|
|
|
|
from app.services import image_upscaler
|
|
from app.models.job import Job
|
|
from app.models.asset import Asset
|
|
from app.config import settings
|
|
|
|
# Mock objects
|
|
mock_db = MagicMock()
|
|
mock_job = MagicMock()
|
|
mock_asset = MagicMock()
|
|
|
|
mock_job.id = "test_job_img_123"
|
|
mock_job.input_asset_ids = ["asset_img_123"]
|
|
# Frontend now sends correct strings, e.g. "High Fidelity V2"
|
|
mock_job.input_data = {
|
|
"scale": 4,
|
|
"model": "High Fidelity V2",
|
|
"sharpen": 0.5, # Frontend sends "sharpen"
|
|
"denoise_strength": 0.4, # Frontend sends denoise_strength
|
|
"face_enhancement": True
|
|
}
|
|
|
|
mock_asset.id = "asset_img_123"
|
|
mock_asset.file_path = "test_image.jpg"
|
|
mock_asset.original_filename = "test_image.jpg"
|
|
mock_asset.mime_type = "image/jpeg"
|
|
mock_asset.width = 1000
|
|
mock_asset.height = 1000
|
|
|
|
# Mock DB queries
|
|
mock_db.query.return_value.filter.return_value.first.side_effect = [mock_job, mock_asset]
|
|
# Mock file read
|
|
file_mock = MagicMock()
|
|
file_mock.__enter__.return_value.read.return_value = b"fake_image_data"
|
|
|
|
async def test_topaz_image_payload():
|
|
print("Testing Topaz Image Upscaling Payload Construction...")
|
|
|
|
with patch("app.services.image_upscaler.SessionLocal", return_value=mock_db):
|
|
with patch("builtins.open", file_mock):
|
|
with patch("app.services.image_upscaler.httpx.AsyncClient") as MockClient:
|
|
mock_client_instance = MockClient.return_value.__aenter__.return_value
|
|
|
|
# Mock API responses
|
|
mock_client_instance.post.return_value.status_code = 200
|
|
mock_client_instance.post.return_value.json.return_value = {"id": "req_123"}
|
|
mock_client_instance.get.return_value.json.return_value = {"status": "running"}
|
|
# We expect it to loop/fail on polling or download, but we check the POST
|
|
|
|
try:
|
|
# Run usage (will likely error on polling loop or sleep, but POST happens first)
|
|
# We accept error after POST
|
|
try:
|
|
await image_upscaler.upscale("test_job_img_123")
|
|
except Exception as e:
|
|
print(f" [INFO] Execution stopped as expected: {e}")
|
|
|
|
except Exception as e:
|
|
print(f" [ERROR] {e}")
|
|
|
|
# VERIFY POST CALL
|
|
# Calling endpoint: https://api.topazlabs.com/image/v1/enhance/async
|
|
found_call = False
|
|
for call in mock_client_instance.post.call_args_list:
|
|
args, kwargs = call
|
|
url = args[0]
|
|
if "enhance/async" in url:
|
|
found_call = True
|
|
print(" [SUCCESS] API Endpoint Correct (enhance/async)")
|
|
|
|
data = kwargs.get("data", {})
|
|
|
|
print(f" Payload Data: {data}")
|
|
|
|
# Verify mappings
|
|
# sharpen -> sharpen
|
|
# denoise_strength -> denoise
|
|
# model -> model
|
|
|
|
if data.get("model") == "High Fidelity V2":
|
|
print(" [VERIFIED] Model name correctly passed")
|
|
else:
|
|
print(f" [FAILED] Model name mismatch. Got: {data.get('model')}")
|
|
|
|
if str(data.get("sharpen")) == "0.5":
|
|
print(" [VERIFIED] Sharpen parameter correctly mapped")
|
|
else:
|
|
print(f" [FAILED] Sharpen mismatch. Got: {data.get('sharpen')}")
|
|
|
|
if str(data.get("denoise")) == "0.4":
|
|
print(" [VERIFIED] Denoise parameter correctly mapped from denoise_strength")
|
|
else:
|
|
print(f" [FAILED] Denoise mismatch. Got: {data.get('denoise')}")
|
|
|
|
if data.get("face_enhancement") == "true":
|
|
print(" [VERIFIED] Face Enhancement enabled")
|
|
else:
|
|
print(f" [FAILED] Face Enhancement missing or false")
|
|
|
|
if not found_call:
|
|
print(" [FAILED] POST to enhance/async not found")
|
|
|
|
if __name__ == "__main__":
|
|
asyncio.run(test_topaz_image_payload())
|