115 lines
4.6 KiB
Python
115 lines
4.6 KiB
Python
import os
|
|
import sys
|
|
import time
|
|
import unittest
|
|
from fastapi.testclient import TestClient
|
|
|
|
# Ensure tour-comfy is in the import path
|
|
_HERE = os.path.dirname(os.path.abspath(__file__))
|
|
if _HERE not in sys.path:
|
|
sys.path.append(_HERE)
|
|
|
|
from edit_api import app, _load_output_dir
|
|
|
|
class TestQwenBackendRegression(unittest.TestCase):
|
|
@classmethod
|
|
def setUpClass(cls):
|
|
cls.client = TestClient(app)
|
|
cls.output_dir = _load_output_dir()
|
|
|
|
# Real-data filenames specified in the issue description
|
|
cls.real_img_1 = "20260618_053519_1_20260618_053458_image.png"
|
|
cls.real_img_2 = "20260618_052537_0_20260618_052526_image.png"
|
|
cls.wireframe_img = "up_20260628_104641_image.png"
|
|
|
|
# Validate that they exist on disk to run real end-to-end tests
|
|
cls.run_real_tests = True
|
|
img1_path = os.path.join(cls.output_dir, cls.real_img_1)
|
|
img2_path = os.path.join(cls.output_dir, cls.real_img_2)
|
|
|
|
if not os.path.exists(img1_path) or not os.path.exists(img2_path):
|
|
cls.run_real_tests = False
|
|
print(f"[test] Real images not found at {img1_path} or {img2_path}. Running tests with fallback/mock assertion modes.")
|
|
|
|
def test_multi_ref_qwen_endpoint(self):
|
|
"""Test multi-reference generation endpoint with specified real-data images."""
|
|
if not self.run_real_tests:
|
|
self.skipTest("Skipping real backend test due to missing real image artifacts on this node.")
|
|
|
|
payload = {
|
|
"filenames": [self.real_img_1, self.real_img_2],
|
|
"prompt": "standing girl looking at camera, high quality, highly detailed",
|
|
"seed": 42,
|
|
"max_area": 512 * 512,
|
|
"pad_outpaint": False
|
|
}
|
|
|
|
response = self.client.post("/multi-ref", json=payload)
|
|
self.assertEqual(response.status_code, 200)
|
|
|
|
data = response.json()
|
|
self.assertIn("job_id", data)
|
|
job_id = data["job_id"]
|
|
|
|
# Poll status until done (limit to 300 seconds to prevent hanging if backend is offline)
|
|
print(f"[test] Polling multi-ref job {job_id} until completion...")
|
|
start_time = time.time()
|
|
completed = False
|
|
|
|
while time.time() - start_time < 300:
|
|
status_resp = self.client.get(f"/batch/{job_id}")
|
|
if status_resp.status_code == 200:
|
|
job_data = status_resp.json()
|
|
status = job_data.get("status")
|
|
print(f"[test] Job status: {status} (done: {job_data.get('done')}/{job_data.get('total')})")
|
|
if status == "done":
|
|
completed = True
|
|
break
|
|
elif status in ["error", "cancelled"]:
|
|
break
|
|
time.sleep(3)
|
|
|
|
self.assertTrue(completed, "Multi-ref generation job did not complete successfully in time.")
|
|
|
|
def test_scenery_qwen_endpoint(self):
|
|
"""Test scenery generation endpoint with specified wireframe scene image and prompt."""
|
|
if not self.run_real_tests:
|
|
self.skipTest("Skipping real backend test due to missing real image artifacts on this node.")
|
|
|
|
payload = {
|
|
"model_filename": self.real_img_1,
|
|
"scene_image": self.wireframe_img,
|
|
"prompt": "replace person from image 1 naturally with the person from Image 2. Keep the exact position of Image 1, and the exact person of Image 2",
|
|
"seed": 42
|
|
}
|
|
|
|
response = self.client.post("/generate-scenery", json=payload)
|
|
self.assertEqual(response.status_code, 200)
|
|
|
|
data = response.json()
|
|
self.assertIn("job_id", data)
|
|
job_id = data["job_id"]
|
|
|
|
# Poll status until done
|
|
print(f"[test] Polling scenery job {job_id} until completion...")
|
|
start_time = time.time()
|
|
completed = False
|
|
|
|
while time.time() - start_time < 300:
|
|
status_resp = self.client.get(f"/batch/{job_id}")
|
|
if status_resp.status_code == 200:
|
|
job_data = status_resp.json()
|
|
status = job_data.get("status")
|
|
print(f"[test] Job status: {status} (done: {job_data.get('done')}/{job_data.get('total')})")
|
|
if status == "done":
|
|
completed = True
|
|
break
|
|
elif status in ["error", "cancelled"]:
|
|
break
|
|
time.sleep(3)
|
|
|
|
self.assertTrue(completed, "Scenery generation job did not complete successfully in time.")
|
|
|
|
if __name__ == "__main__":
|
|
unittest.main()
|