import os import sys import json import shutil import unittest from datetime import datetime as _dt from PIL import Image # Ensure tour-comfy is in import path so we can import the FastAPI app and database module sys.path.append(os.path.dirname(os.path.abspath(__file__))) from fastapi.testclient import TestClient import database from edit_api import app, _load_output_dir class TestAPIRegression(unittest.TestCase): @classmethod def setUpClass(cls): # Determine output directory cls.output_dir = _load_output_dir() cls.client = TestClient(app) # Unique mock identifiers to avoid conflicts cls.test_ref_filename = "test_regression_ref_image_123.png" cls.test_other_filename = "test_regression_other_image_123.png" cls.group_id = "test_regression_group_123" cls.test_ref_path = os.path.join(cls.output_dir, cls.test_ref_filename) cls.test_other_path = os.path.join(cls.output_dir, cls.test_other_filename) # Ensure any leftover artifacts from past failed runs are cleaned cls._cleanup_database() cls._cleanup_files() # Create dummy test files (100x100 pixels, RGB format) img = Image.new("RGB", (100, 100), color=(255, 0, 0)) img.save(cls.test_ref_path, "PNG") img_other = Image.new("RGB", (100, 100), color=(0, 255, 0)) img_other.save(cls.test_other_path, "PNG") # Mock metadata cls.name = "Regression Test Character" cls.tags = ["VISIBLE", "LIKE", "21+"] cls.embedding = [0.1] * 1024 cls.clip_description = "A dummy regression test image description" cls.prompt = "masterpiece, high quality" cls.pose = "standing" cls.group_name = "Regression Group" cls.hidden = False cls.has_background = True cls.has_clothing = False cls.is_source = True cls.pose_description = "The model is standing and looking directly at the camera." cls.pose_skeleton = '{"keypoints": [1, 2, 3]}' # Insert reference image into the database database.upsert_person( cls.test_ref_filename, filepath=cls.test_ref_path, name=cls.name, group_id=cls.group_id, tags=cls.tags, embedding=cls.embedding, clip_description=cls.clip_description, prompt=cls.prompt, pose=cls.pose, sort_order=0, group_name=cls.group_name, hidden=cls.hidden, has_background=cls.has_background, has_clothing=cls.has_clothing, is_source=cls.is_source, pose_description=cls.pose_description, pose_skeleton=cls.pose_skeleton ) # Insert second image in same group for reordering test database.upsert_person( cls.test_other_filename, filepath=cls.test_other_path, name=cls.name, group_id=cls.group_id, tags=cls.tags, embedding=cls.embedding, clip_description=cls.clip_description, prompt=cls.prompt, pose=cls.pose, sort_order=1, group_name=cls.group_name, hidden=cls.hidden, has_background=cls.has_background, has_clothing=cls.has_clothing, is_source=cls.is_source, pose_description=cls.pose_description, pose_skeleton=cls.pose_skeleton ) # Track dynamically created files and database rows for cleanup cls.created_files = [cls.test_ref_path, cls.test_other_path] cls.created_db_rows = [cls.test_ref_filename, cls.test_other_filename] @classmethod def tearDownClass(cls): # Cleanup cls._cleanup_database() cls._cleanup_files() @classmethod def _cleanup_database(cls): # Safely delete any inserted test rows from the person database table conn = database.get_db_connection() cur = conn.cursor() try: # Delete any filenames starting with test_regression or containing ts_crop/ts_pad cur.execute(""" DELETE FROM person WHERE filename LIKE 'test_regression_%%' OR filename LIKE '%%_crop_test_regression_%%' OR filename LIKE '%%_pad_test_regression_%%' OR group_id = %s """, (cls.group_id,)) conn.commit() except Exception as e: print(f"Error cleaning up database: {e}") conn.rollback() finally: cur.close() database._put_db_connection(conn) @classmethod def _cleanup_files(cls): # Clean up files matching test patterns for f in os.listdir(cls.output_dir): if "test_regression_" in f or "_crop_" in f or "_pad_" in f: p = os.path.join(cls.output_dir, f) try: if os.path.exists(p): os.remove(p) except Exception as e: print(f"Error removing file {p}: {e}") def assertEmbeddingEqual(self, val1, val2): if isinstance(val1, str): val1 = [float(x) for x in val1.strip("[]").split(",")] if isinstance(val2, str): val2 = [float(x) for x in val2.strip("[]").split(",")] self.assertEqual(len(val1), len(val2)) for x, y in zip(val1, val2): self.assertAlmostEqual(x, y, places=4) def test_01_duplicate_copies_all_pose_and_meta_details(self): # Send duplicate request to FastAPI response = self.client.post(f"/images/{self.test_ref_filename}/duplicate") self.assertEqual(response.status_code, 200) res_data = response.json() self.assertEqual(res_data["status"], "success") new_filename = res_data["new_filename"] self.created_db_rows.append(new_filename) new_path = os.path.join(self.output_dir, os.path.basename(new_filename)) self.created_files.append(new_path) # Assert duplicate file exists on disk self.assertTrue(os.path.exists(new_path), f"Duplicated file {new_path} not found on disk") # Assert all metadata has been accurately duplicated in DB person = database.get_person(new_filename) self.assertIsNotNone(person, "Duplicated database entry not found") # Column mappings as in database.py get_person: # SELECT name, group_id, tags, embedding, clip_description, filepath, # prompt, pose, sort_order, group_name, hidden, has_background, source_refs, # has_clothing, is_source, pose_description, pose_skeleton self.assertEqual(person[0], self.name) self.assertEqual(person[1], self.group_id) # Tags (assert LIKE and 21+ are preserved) tags_list = json.loads(person[2]) if isinstance(person[2], str) else person[2] self.assertIn("LIKE", tags_list) self.assertIn("21+", tags_list) self.assertIn("VISIBLE", tags_list) # Verify 21+ is just a standard tag, not a safety blocker self.assertIn("21+", tags_list) self.assertEmbeddingEqual(person[3], self.embedding) self.assertEqual(person[4], self.clip_description) self.assertEqual(person[5], new_path) self.assertEqual(person[6], self.prompt) self.assertEqual(person[7], self.pose) self.assertEqual(person[9], self.group_name) self.assertEqual(person[10], self.hidden) self.assertEqual(person[11], self.has_background) # source_refs should refer to original source_refs = json.loads(person[12]) if isinstance(person[12], str) else person[12] self.assertIn(self.test_ref_filename, source_refs) self.assertEqual(person[13], self.has_clothing) self.assertEqual(person[14], self.is_source) self.assertEqual(person[15], self.pose_description) self.assertEqual(person[16], self.pose_skeleton) def test_02_crop_as_copy_copies_all_pose_and_meta_details(self): # Crop region: (10, 10, 90, 90), as_copy=True req_payload = { "x1": 10, "y1": 10, "x2": 90, "y2": 90, "as_copy": True } response = self.client.post(f"/images/{self.test_ref_filename}/crop", json=req_payload) self.assertEqual(response.status_code, 200) res_data = response.json() self.assertEqual(res_data["status"], "success") new_filename = res_data["new_filename"] self.created_db_rows.append(new_filename) new_path = os.path.join(self.output_dir, os.path.basename(new_filename)) self.created_files.append(new_path) # Verify physical file existence and crop dimensions (should be 80x80) self.assertTrue(os.path.exists(new_path), f"Cropped file {new_path} not found on disk") cropped_img = Image.open(new_path) self.assertEqual(cropped_img.size, (80, 80)) # Verify database entry has complete metadata person = database.get_person(new_filename) self.assertIsNotNone(person, "Cropped copy database entry not found") self.assertEqual(person[0], self.name) self.assertEqual(person[1], self.group_id) tags_list = json.loads(person[2]) if isinstance(person[2], str) else person[2] self.assertIn("LIKE", tags_list) self.assertIn("21+", tags_list) self.assertIn("VISIBLE", tags_list) self.assertEmbeddingEqual(person[3], self.embedding) self.assertEqual(person[4], self.clip_description) self.assertEqual(person[5], new_path) self.assertEqual(person[6], self.prompt) self.assertEqual(person[7], self.pose) self.assertEqual(person[9], self.group_name) self.assertEqual(person[10], self.hidden) self.assertEqual(person[11], self.has_background) source_refs = json.loads(person[12]) if isinstance(person[12], str) else person[12] self.assertIn(self.test_ref_filename, source_refs) self.assertEqual(person[13], self.has_clothing) self.assertEqual(person[14], self.is_source) self.assertEqual(person[15], self.pose_description) self.assertEqual(person[16], self.pose_skeleton) def test_03_crop_in_place_keeps_all_meta_information(self): # Crop region: (20, 20, 80, 80), as_copy=False (in-place) req_payload = { "x1": 20, "y1": 20, "x2": 80, "y2": 80, "as_copy": False } # First verify original size is 100x100 orig_img = Image.open(self.test_ref_path) self.assertEqual(orig_img.size, (100, 100)) response = self.client.post(f"/images/{self.test_ref_filename}/crop", json=req_payload) self.assertEqual(response.status_code, 200) res_data = response.json() self.assertEqual(res_data["status"], "success") self.assertEqual(res_data["new_filename"], self.test_ref_filename) # Verify physical file size updated to 60x60 cropped_img = Image.open(self.test_ref_path) self.assertEqual(cropped_img.size, (60, 60)) # Verify database entry has complete metadata untouched person = database.get_person(self.test_ref_filename) self.assertIsNotNone(person, "Database entry not found after in-place crop") self.assertEqual(person[0], self.name) self.assertEqual(person[1], self.group_id) tags_list = json.loads(person[2]) if isinstance(person[2], str) else person[2] self.assertIn("LIKE", tags_list) self.assertIn("21+", tags_list) self.assertEmbeddingEqual(person[3], self.embedding) self.assertEqual(person[4], self.clip_description) self.assertEqual(person[6], self.prompt) self.assertEqual(person[7], self.pose) self.assertEqual(person[9], self.group_name) self.assertEqual(person[10], self.hidden) self.assertEqual(person[11], self.has_background) self.assertEqual(person[13], self.has_clothing) self.assertEqual(person[14], self.is_source) self.assertEqual(person[15], self.pose_description) self.assertEqual(person[16], self.pose_skeleton) def test_04_pad_as_copy_copies_all_pose_and_meta_details(self): # Pad canvas: expand top and bottom by 10 pixels, as_copy=True req_payload = { "top": 10, "right": 0, "bottom": 10, "left": 0, "as_copy": True, "fill": "transparent", "outpaint": False } # Original size at this point is 60x60 (due to test_03 crop) response = self.client.post(f"/images/{self.test_ref_filename}/pad", json=req_payload) self.assertEqual(response.status_code, 200) res_data = response.json() self.assertEqual(res_data["status"], "success") new_filename = res_data["new_filename"] self.created_db_rows.append(new_filename) new_path = os.path.join(self.output_dir, os.path.basename(new_filename)) self.created_files.append(new_path) # Verify physical file existence and pad dimensions (should be 60x80) self.assertTrue(os.path.exists(new_path), f"Padded file {new_path} not found on disk") padded_img = Image.open(new_path) self.assertEqual(padded_img.size, (60, 80)) # Verify database entry has complete metadata person = database.get_person(new_filename) self.assertIsNotNone(person, "Padded copy database entry not found") self.assertEqual(person[0], self.name) self.assertEqual(person[1], self.group_id) tags_list = json.loads(person[2]) if isinstance(person[2], str) else person[2] self.assertIn("LIKE", tags_list) self.assertIn("21+", tags_list) self.assertIn("VISIBLE", tags_list) self.assertEmbeddingEqual(person[3], self.embedding) self.assertEqual(person[4], self.clip_description) self.assertEqual(person[5], new_path) self.assertEqual(person[6], self.prompt) self.assertEqual(person[7], self.pose) self.assertEqual(person[9], self.group_name) self.assertEqual(person[10], self.hidden) self.assertEqual(person[11], self.has_background) source_refs = json.loads(person[12]) if isinstance(person[12], str) else person[12] self.assertIn(self.test_ref_filename, source_refs) self.assertEqual(person[13], self.has_clothing) self.assertEqual(person[14], self.is_source) self.assertEqual(person[15], self.pose_description) self.assertEqual(person[16], self.pose_skeleton) def test_05_reorder_group_updates_sort_orders(self): # Verify original order order_rows_before = database.get_group_order(self.group_id) filenames_before = [r[0] for r in order_rows_before] self.assertIn(self.test_ref_filename, filenames_before) self.assertIn(self.test_other_filename, filenames_before) # Reverse the order of files and submit new_order = [self.test_other_filename, self.test_ref_filename] req_payload = { "filenames": new_order } response = self.client.post(f"/groups/{self.group_id}/order", json=req_payload) self.assertEqual(response.status_code, 200) res_data = response.json() self.assertEqual(res_data["group_id"], self.group_id) self.assertEqual(res_data["filenames"], new_order) # Get order from database and verify state reflects the updates order_rows_after = database.get_group_order(self.group_id) filenames_after = [r[0] for r in order_rows_after] # Verify custom sort orders are explicitly 0, 1 self.assertEqual(filenames_after[:2], new_order) person_other = database.get_person(self.test_other_filename) person_ref = database.get_person(self.test_ref_filename) # Column 8 in database.get_person is sort_order self.assertEqual(person_other[8], 0) self.assertEqual(person_ref[8], 1) def test_06_list_images_with_bypass_static(self): # Call with bypass_static=True response = self.client.get("/images?archived=true&bypass_static=true") self.assertEqual(response.status_code, 200) res_data = response.json() self.assertIn("images", res_data) # Find our reference image in the list images = res_data["images"] ref_img = next((img for img in images if img["filename"] == self.test_ref_filename), None) self.assertIsNotNone(ref_img, "Test reference image not found in bypassed images list") # Verify complete mapped metadata fields are correctly returned when bypassing static file self.assertEqual(ref_img["name"], self.name) self.assertEqual(ref_img["group_id"], self.group_id) self.assertEqual(ref_img["prompt"], self.prompt) self.assertEqual(ref_img["pose"], self.pose) self.assertEqual(ref_img["group_name"], self.group_name) self.assertEqual(ref_img["hidden"], self.hidden) self.assertEqual(ref_img["has_background"], self.has_background) self.assertEqual(ref_img["has_clothing"], self.has_clothing) self.assertEqual(ref_img["is_source"], self.is_source) self.assertEqual(ref_img["pose_description"], self.pose_description) # Verify tags list self.assertIn("LIKE", ref_img["tags"]) self.assertIn("21+", ref_img["tags"]) # Verify skeleton self.assertEqual(ref_img["pose_skeleton"], self.pose_skeleton) def test_07_concurrent_reordering_deadlock_resilience(self): import threading import random import time exceptions = [] def worker(thread_idx): try: for i in range(15): # Alternate the order to trigger lock conflicts if i % 2 == 0: order = [self.test_other_filename, self.test_ref_filename] else: order = [self.test_ref_filename, self.test_other_filename] database.set_group_order(self.group_id, order) time.sleep(0.01) except Exception as e: exceptions.append((thread_idx, e)) threads = [] for j in range(5): t = threading.Thread(target=worker, args=(j,)) threads.append(t) t.start() for t in threads: t.join() # If any deadlock or serialization error was uncaught, exceptions won't be empty if exceptions: for tid, err in exceptions: print(f"[test] Thread {tid} encountered error: {err}") self.assertEqual(len(exceptions), 0, f"Encountered concurrent set_group_order errors: {exceptions}") def test_08_deep_metadata_exposure_and_automatic_background_processing(self): # Trigger the /images list response = self.client.get("/images?archived=true&bypass_static=true") self.assertEqual(response.status_code, 200) res_data = response.json() self.assertIn("images", res_data) # Verify that the new keys are present in each image dict for img in res_data["images"]: self.assertIn("people_count", img) self.assertIn("anatomical_completeness", img) self.assertIn("facial_direction", img) self.assertIn("objects", img) # Verify that manually triggering metadata backfill for our ref image returns expected structure backfill_payload = {"filenames": [self.test_ref_filename]} response = self.client.post("/images/backfill-metadata", json=backfill_payload) self.assertEqual(response.status_code, 200) res_backfill = response.json() self.assertEqual(res_backfill["status"], "completed") self.assertEqual(res_backfill["total"], 1) def test_09_upload_group_assignment(self): from unittest.mock import patch import io # 1. Test skip_poses=True, group_id=None # It should assign group_id to the base name of the uploaded file. img_bytes = io.BytesIO() Image.new("RGB", (100, 100), color=(0, 0, 255)).save(img_bytes, format="PNG") img_bytes.seek(0) response = self.client.post( "/upload", files={"image": ("test_upload_image.png", img_bytes, "image/png")}, data={"skip_poses": "true"} ) self.assertEqual(response.status_code, 200) res_data = response.json() self.assertEqual(res_data["status"], "added") self.assertEqual(res_data["group_id"], "test_upload_image.png") self.created_db_rows.append(res_data["filename"]) self.created_files.append(os.path.join(self.output_dir, res_data["filename"])) # 2. Test skip_poses=False, group_id=None (e.g. CTRL+v New group with run poses) # It should NOT collapse to "paste.png" or "test_upload_image.png", but generate a unique up_XXXX group ID! img_bytes2 = io.BytesIO() Image.new("RGB", (100, 100), color=(0, 255, 255)).save(img_bytes2, format="PNG") img_bytes2.seek(0) with patch("edit_api._process_upload") as mock_process: response = self.client.post( "/upload", files={"image": ("test_upload_image2.png", img_bytes2, "image/png")}, data={"skip_poses": "false"} ) self.assertEqual(response.status_code, 200) res_data2 = response.json() self.assertEqual(res_data2["status"], "processing") # It should generate a unique ID starting with up_ self.assertTrue(res_data2["group_id"].startswith("up_")) self.assertNotEqual(res_data2["group_id"], "test_upload_image2.png") # Verify background task was called with the unique group_id mock_process.assert_called_once() args, kwargs = mock_process.call_args # group_id is the 5th positional arg or a kwarg called_gid = kwargs.get("group_id") or args[4] self.assertTrue(called_gid.startswith("up_")) self.created_files.append(os.path.join(self.output_dir, res_data2["filename"])) def test_10_idle_backfill_non_image_filtering(self): from unittest.mock import patch, MagicMock import edit_api # Reset failed backfill filenames set before testing edit_api._failed_backfill_filenames.clear() # Mock database.list_persons to return rows with: # 1. A video file # 2. A non-image file # 3. An image that does exist on disk but we want to backfill # 4. A normal image where people_count is already present mock_persons = [ ("video.mp4", "Video", "g1", None, "", None, 0, "", False, False, None, False, "video", None, False, False, [], None, None, None, None, None, []), ("notes.txt", "Text", "g1", None, "", None, 1, "", False, False, None, False, "image", None, False, False, [], None, None, None, None, None, []), ("processed_image.png", "Image1", "g1", None, "", None, 2, "", False, False, None, False, "image", None, False, False, [], None, None, 1, "Full", "Front", []), ("test_ref_image_123.png", "Image2", "g1", None, "", None, 3, "", False, False, None, False, "image", None, False, False, [], None, None, None, None, None, []), ] # First, test the fast-fail guard inside _process_image_for_metadata directly (not mocked yet) res_mp4 = edit_api._process_image_for_metadata("video.mp4") self.assertIsNone(res_mp4) res_txt = edit_api._process_image_for_metadata("notes.txt") self.assertIsNone(res_txt) # Let's mock os.path.exists to return True for any path we check with patch("os.path.exists", return_value=True), \ patch("database.list_persons", return_value=mock_persons), \ patch("edit_api._process_image_for_metadata") as mock_process: # Second, let's replicate/test the candidate selection loop inside edit_api._idle_turntable_daemon output_dir = "/mnt/zim/tour-comfy/output" _failed_backfill_filenames = edit_api._failed_backfill_filenames persons = mock_persons legacy_candidate = None for row in persons: fname = row[0] content_type = row[12] if len(row) > 12 else None people_count = row[19] if len(row) > 19 else None if fname.startswith("_turntable/"): continue if content_type == 'video': continue if not fname.lower().endswith(('.png', '.jpg', '.jpeg', '.webp')): continue fpath = os.path.join(output_dir, fname) if not os.path.exists(fpath): continue if people_count is None and fname not in _failed_backfill_filenames: legacy_candidate = fname break self.assertEqual(legacy_candidate, "test_ref_image_123.png") # Let's test failed backfill insertion if _process_image_for_metadata returns None mock_process.return_value = None # Simulating the idle turntable loop processing the candidate if legacy_candidate: try: res = edit_api._process_image_for_metadata(legacy_candidate) if res is None: _failed_backfill_filenames.add(legacy_candidate) except Exception as ex: _failed_backfill_filenames.add(legacy_candidate) self.assertIn("test_ref_image_123.png", _failed_backfill_filenames) def test_11_video_listing_refresh_and_trim_synchronization(self): from unittest.mock import patch import edit_api # Mock wireframe folder and output directories with patch("edit_api._load_wireframe_dir", return_value="/tmp/test_wireframe"), \ patch("edit_api._load_output_dir", return_value=self.output_dir), \ patch("os.path.isdir", return_value=True), \ patch("os.listdir", return_value=["dance.mp4", "dance_10s-20s.mp4", "notes.txt"]): # Verify that calling /videos?refresh=true triggers grouping correctly response = self.client.get("/videos?refresh=true") self.assertEqual(response.status_code, 200) data = response.json() # Grouping checks self.assertIn("groups", data) groups = data["groups"] # Find group with stem "dance" dance_grp = next((g for g in groups if g["stem"] == "dance"), None) self.assertIsNotNone(dance_grp) self.assertEqual(dance_grp["video"], "dance.mp4") self.assertIn("dance_10s-20s.mp4", dance_grp["clips"]) def test_12_standalone_scenery_sorting_and_existence_check(self): from unittest.mock import patch import edit_api import database # Insert some mock scenery generation records so we can test last_used sorting conn = database.get_db_connection() cur = conn.cursor() try: cur.execute(""" INSERT INTO person (filename, group_id, source_refs) VALUES ('20260629_130000_sc_test.png', 'g1', '["test_person.png", "wireframe:used_scenery_2.png"]') """) cur.execute(""" INSERT INTO person (filename, group_id, source_refs) VALUES ('20260629_130500_sc_test.png', 'g1', '["test_person.png", "wireframe:used_scenery_1.png"]') """) conn.commit() finally: cur.close() database._put_db_connection(conn) try: # Mock wireframe folder and output directories with patch("edit_api._load_wireframe_dir", return_value="/tmp/test_wireframe"), \ patch("edit_api._load_output_dir", return_value=self.output_dir), \ patch("os.path.isdir", return_value=True), \ patch("os.path.exists", return_value=True), \ patch("os.path.getmtime", return_value=12345.0), \ patch("os.listdir", return_value=["used_scenery_1.png", "used_scenery_2.png", "unused_scenery_3.png"]): response = self.client.get("/videos?refresh=true") self.assertEqual(response.status_code, 200) data = response.json() # Check standalone images exist in the response self.assertIn("standalone_images", data) standalone = data["standalone_images"] # Should contain our wireframes sorted with last used on top: # 'used_scenery_1.png' was used in the most recent scenery generation (at 13:05:00) # 'used_scenery_2.png' was used in the older scenery generation (at 13:00:00) # 'unused_scenery_3.png' was not used self.assertEqual(standalone[0], "used_scenery_1.png") self.assertEqual(standalone[1], "used_scenery_2.png") self.assertEqual(standalone[2], "unused_scenery_3.png") # Also verify that /scenery/library filters out missing files # Let's mock _get_cached_file_meta: '20260629_130500_sc_test.png' exists, but '20260629_130000_sc_test.png' is missing def mock_cached_meta(fname, out_dir): if fname == "20260629_130500_sc_test.png": return True, 12345.0 return False, 0.0 with patch("edit_api._get_cached_file_meta", side_effect=mock_cached_meta): response = self.client.get("/scenery/library") self.assertEqual(response.status_code, 200) lib_data = response.json() # Verify that only the existing scenery item '20260629_130500_sc_test.png' is returned all_filenames = [] for g in lib_data["groups"]: for item in g["items"]: all_filenames.append(item["filename"]) for item in lib_data["ungrouped"]: all_filenames.append(item["filename"]) self.assertIn("20260629_130500_sc_test.png", all_filenames) self.assertNotIn("20260629_130000_sc_test.png", all_filenames) finally: # Clean up test DB rows conn = database.get_db_connection() cur = conn.cursor() try: cur.execute("DELETE FROM person WHERE filename IN ('20260629_130000_sc_test.png', '20260629_130500_sc_test.png')") conn.commit() finally: cur.close() database._put_db_connection(conn) def test_13_scenery_metadata_extraction_and_refs(self): from unittest.mock import patch, MagicMock import edit_api import database import json import os from PIL import Image # Write dummy model image on disk model_filename = "model_image_123.png" model_path = os.path.join(self.output_dir, model_filename) Image.new("RGB", (10, 10), color="red").save(model_path) # Set up a test person record in DB for the model conn = database.get_db_connection() cur = conn.cursor() try: cur.execute(""" INSERT INTO person (filename, group_id) VALUES (%s, 'model_group_123') """, (model_filename,)) conn.commit() finally: cur.close() database._put_db_connection(conn) try: # Let's mock _run_pipeline, generation of embeddings, etc. # We want to verify that _scenery_worker correctly structures 'source_refs' # and triggers metadata extraction. mock_png = b"dummy_png_bytes" edit_api.jobs["test_scenery_job"] = {"status": "running"} with patch("edit_api._run_pipeline", return_value=mock_png), \ patch("edit_api._load_output_dir", return_value=self.output_dir), \ patch("edit_api.embeddings.generate_embedding", return_value=None), \ patch("edit_api._metadata_executor.submit") as mock_submit: # Run the scenery worker synchronously for the test dummy_scene = Image.new("RGB", (100, 100), color="blue") edit_api._scenery_worker( job_id="test_scenery_job", model_filename=model_filename, scene_pil=dummy_scene, prompt="test scenery prompt", seed=42, extra_pils=[], scene_video="test_video.mp4", scene_image="test_image.png", extra_filename="test_extra.png" ) # Find the newly created scenery filename scenery_filename = edit_api.jobs["test_scenery_job"]["output"] self.assertIsNotNone(scenery_filename) # Check database record to verify refs are correct person = database.get_person(scenery_filename) self.assertIsNotNone(person) self.assertEqual(person[1], "model_group_123") # Correctly pulled group_id refs = json.loads(person[12]) # source_refs is at index 12 in tuple or get_person response self.assertIn(model_filename, refs) self.assertIn("video:test_video.mp4", refs) self.assertIn("wireframe:test_image.png", refs) self.assertIn("test_extra.png", refs) # Verify that _metadata_executor.submit was called to trigger background metadata extraction mock_submit.assert_called_once_with(edit_api._process_image_for_metadata, scenery_filename) finally: # Clean up all created files and database entries if os.path.exists(model_path): os.remove(model_path) conn = database.get_db_connection() cur = conn.cursor() try: cur.execute("DELETE FROM person WHERE filename = 'model_image_123.png'") cur.execute("DELETE FROM person WHERE filename LIKE '%_sc_model_image_123.png'") conn.commit() finally: cur.close() database._put_db_connection(conn) # Clean up generated scenery file on disk try: scenery_filename = edit_api.jobs.get("test_scenery_job", {}).get("output") if scenery_filename: out_path = os.path.join(self.output_dir, scenery_filename) if os.path.exists(out_path): os.remove(out_path) except Exception: pass def test_13b_scenery_validation_checks(self): from unittest.mock import patch import edit_api from fastapi import HTTPException # Test 1: model image is completely missing -> raise 404 req_missing = edit_api.SceneryRequest( model_filename="non_existent_model_file_999.png", scene_image="some_scene.png", prompt="replace person from image 1 with the person from Image 2" ) with self.assertRaises(HTTPException) as ctx: edit_api.generate_scenery(req_missing) self.assertEqual(ctx.exception.status_code, 404) self.assertIn("Model image not found", ctx.exception.detail) # Write dummy model image on disk for subsequent tests model_filename = "test_model_validation_123.png" model_path = os.path.join(self.output_dir, model_filename) Image.new("RGB", (10, 10), color="red").save(model_path) try: # Test 2: model and scene images are the same file -> raise 400 req_identical = edit_api.SceneryRequest( model_filename=model_filename, scene_image=model_filename, prompt="replace person from image 1 with the person from Image 2" ) with self.assertRaises(HTTPException) as ctx: edit_api.generate_scenery(req_identical) self.assertEqual(ctx.exception.status_code, 400) self.assertIn("cannot be the same image", ctx.exception.detail) # Test 3: prompt refinement "image 1/2" -> "Picture 1/2" # We mock the thread start and save_db_prompt to verify the processed prompt saved_prompt = [] def mock_save_db_prompt(ptype, prompt_text, meta): saved_prompt.append(prompt_text) dummy_scene_filename = "dummy_scene_validation_123.png" dummy_scene_path = os.path.join(self.output_dir, dummy_scene_filename) Image.new("RGB", (10, 10), color="blue").save(dummy_scene_path) try: req_refine = edit_api.SceneryRequest( model_filename=model_filename, scene_image=dummy_scene_filename, prompt="replace person from image 1 with the person from Image 2" ) with patch("edit_api._load_wireframe_dir", return_value=self.output_dir), \ patch("edit_api.database.save_db_prompt", side_effect=mock_save_db_prompt), \ patch("threading.Thread.start") as mock_thread_start: edit_api.generate_scenery(req_refine) self.assertEqual(len(saved_prompt), 1) self.assertEqual(saved_prompt[0], "replace person from Picture 1 with the person from Picture 2") finally: if os.path.exists(dummy_scene_path): os.remove(dummy_scene_path) finally: if os.path.exists(model_path): os.remove(model_path) def test_13c_scenery_generation_insufficient_references_checks(self): import edit_api from fastapi import HTTPException from PIL import Image # Write dummy model image on disk model_filename = "test_scenery_insufficient_model.png" model_path = os.path.join(self.output_dir, model_filename) Image.new("RGB", (10, 10), color="red").save(model_path) try: # Test 1: both scene_image and scene_video are missing -> raise HTTPException 400 req_no_bg = edit_api.SceneryRequest( model_filename=model_filename, prompt="some prompt" ) with self.assertRaises(HTTPException) as ctx: edit_api.generate_scenery(req_no_bg) self.assertEqual(ctx.exception.status_code, 400) self.assertIn("requires a background scene reference", ctx.exception.detail) # Test 2: inside _scenery_worker, if len(refs) < 2, it raises ValueError and sets status to error edit_api.jobs["test_scenery_insufficient_job"] = {"status": "running"} dummy_scene = Image.new("RGB", (10, 10), color="blue") edit_api._scenery_worker( job_id="test_scenery_insufficient_job", model_filename=model_filename, scene_pil=dummy_scene, prompt="test scenery prompt", seed=42, extra_pils=[], scene_video=None, # explicitly None to trigger the < 2 refs check scene_image=None, # explicitly None to trigger the < 2 refs check extra_filename=None ) job = edit_api.jobs["test_scenery_insufficient_job"] self.assertEqual(job["status"], "error") self.assertIn("expected at least 2 source references", job["error"]) finally: if os.path.exists(model_path): os.remove(model_path) def test_14_static_serialization_of_deep_metadata(self): import edit_api import database import json import os from PIL import Image # Create dummy image and person record fn = "test_deep_meta_serialization_123.png" fpath = os.path.join(self.output_dir, fn) Image.new("RGB", (10, 10)).save(fpath) conn = database.get_db_connection() cur = conn.cursor() try: cur.execute(""" INSERT INTO person (filename, filepath, group_id, people_count, anatomical_completeness, facial_direction, objects) VALUES (%s, %s, %s, %s, %s, %s, %s) """, (fn, fpath, "test_group", 2, False, "front", json.dumps([{"tag": "hat", "score": 0.9}]))) conn.commit() finally: cur.close() database._put_db_connection(conn) try: # Trigger write static images.json edit_api._write_all_static() # Read static file and check that the deep metadata properties are serialized correctly static_file = os.path.join(self.output_dir, "_data", "images.json") self.assertTrue(os.path.exists(static_file)) with open(static_file, "r") as f: data = json.load(f) imgs = data.get("images", []) matched = [x for x in imgs if x["filename"] == fn] self.assertEqual(len(matched), 1) item = matched[0] # Assert that the serialized dictionary has all of the deep metadata fields self.assertEqual(item.get("people_count"), 2) self.assertEqual(item.get("anatomical_completeness"), False) self.assertEqual(item.get("facial_direction"), "front") self.assertEqual(item.get("objects"), [{"tag": "hat", "score": 0.9}]) finally: # Clean up if os.path.exists(fpath): os.remove(fpath) conn = database.get_db_connection() cur = conn.cursor() try: cur.execute("DELETE FROM person WHERE filename = %s", (fn,)) conn.commit() finally: cur.close() database._put_db_connection(conn) # Re-generate static to clean up the dummy entry edit_api._write_all_static() def test_15_anatomical_completeness_refinement(self): import edit_api # 17 keypoints matching COCO-17 format: [x, y, score] # Nose (0), shoulders (5,6), elbows (7,8), wrists (9,10), hips (11,12), knees (13,14), ankles (15,16) # All of them are visible with score >= 0.3, and well within bounds [0.1, 0.9] complete_kpts = [ [50, 10, 0.9], # Nose (0) [48, 11, 0.9], # L_Eye (1) [52, 11, 0.9], # R_Eye (2) [46, 12, 0.9], # L_Ear (3) [54, 12, 0.9], # R_Ear (4) [40, 20, 0.9], # L_Shoulder (5) [60, 20, 0.9], # R_Shoulder (6) [35, 35, 0.9], # L_Elbow (7) [65, 35, 0.9], # R_Elbow (8) [30, 50, 0.9], # L_Wrist (9) [70, 50, 0.9], # R_Wrist (10) [42, 55, 0.9], # L_Hip (11) [58, 55, 0.9], # R_Hip (12) [40, 75, 0.9], # L_Knee (13) [60, 75, 0.9], # R_Knee (14) [40, 90, 0.9], # L_Ankle (15) [60, 90, 0.9], # R_Ankle (16) ] # Test complete pose res1 = edit_api._detect_anatomical_completeness(complete_kpts, 100, 100) self.assertTrue(res1) # Test partial pose: missing knees and ankles partial_kpts = [list(kp) for kp in complete_kpts] for idx in [13, 14, 15, 16]: partial_kpts[idx][2] = 0.1 # Hide them res2 = edit_api._detect_anatomical_completeness(partial_kpts, 100, 100) self.assertFalse(res2) # Test cropped pose: ankles too close to bottom edge (y=99 in height=100) cropped_bottom_kpts = [list(kp) for kp in complete_kpts] cropped_bottom_kpts[15][1] = 99 cropped_bottom_kpts[16][1] = 99 res3 = edit_api._detect_anatomical_completeness(cropped_bottom_kpts, 100, 100) self.assertFalse(res3) def test_16_object_bbox_estimation(self): import edit_api # Define some realistic keypoints kpts = [ [50, 10, 0.9], # Nose (0) [48, 11, 0.9], # L_Eye (1) [52, 11, 0.9], # R_Eye (2) [46, 12, 0.9], # L_Ear (3) [54, 12, 0.9], # R_Ear (4) [40, 20, 0.9], # L_Shoulder (5) [60, 20, 0.9], # R_Shoulder (6) [35, 35, 0.9], # L_Elbow (7) [65, 35, 0.9], # R_Elbow (8) [30, 50, 0.9], # L_Wrist (9) [70, 50, 0.9], # R_Wrist (10) [42, 55, 0.9], # L_Hip (11) [58, 55, 0.9], # R_Hip (12) [40, 75, 0.9], # L_Knee (13) [60, 75, 0.9], # R_Knee (14) [40, 90, 0.9], # L_Ankle (15) [60, 90, 0.9], # R_Ankle (16) ] # Test head / hair tag bbox estimation bbox_hair = edit_api._estimate_bbox_for_tag("long_hair", kpts, 100, 100) self.assertIsNotNone(bbox_hair) self.assertEqual(len(bbox_hair), 4) # Bounding box should surround the head/face region, which is y-centered around 10-12. # So upper y (bbox_hair[1]) should be relatively small (close to 0/top) self.assertLess(bbox_hair[1], 15) # Test chest tag bbox estimation bbox_breasts = edit_api._estimate_bbox_for_tag("breasts", kpts, 100, 100) self.assertIsNotNone(bbox_breasts) self.assertEqual(len(bbox_breasts), 4) # Should be below shoulders (y=20) and above hips (y=55). Center around 30. self.assertGreater(bbox_breasts[1], 15) self.assertLess(bbox_breasts[3], 55) # Test stomach tag bbox estimation bbox_navel = edit_api._estimate_bbox_for_tag("navel", kpts, 100, 100) self.assertIsNotNone(bbox_navel) self.assertEqual(len(bbox_navel), 4) # Should be around the stomach area, i.e. between 35 and 55. self.assertGreater(bbox_navel[1], 25) # Test arms tag bbox estimation bbox_arms = edit_api._estimate_bbox_for_tag("sleeves", kpts, 100, 100) self.assertIsNotNone(bbox_arms) self.assertEqual(len(bbox_arms), 4) # Test legs tag bbox estimation bbox_legs = edit_api._estimate_bbox_for_tag("thighs", kpts, 100, 100) self.assertIsNotNone(bbox_legs) self.assertEqual(len(bbox_legs), 4) def test_17_gaze_and_pose_description_enhancements(self): import edit_api # Test 1: Standard upright front facing pose with direct gaze standing_kpts = [ [50, 12.5, 0.9], # Nose (0) [52, 11, 0.9], # L_Eye (1) [48, 11, 0.9], # R_Eye (2) [54, 12, 0.9], # L_Ear (3) [46, 12, 0.9], # R_Ear (4) [60, 20, 0.9], # L_Shoulder (5) [40, 20, 0.9], # R_Shoulder (6) [65, 35, 0.9], # L_Elbow (7) [35, 35, 0.9], # R_Elbow (8) [70, 50, 0.9], # L_Wrist (9) [30, 50, 0.9], # R_Wrist (10) [58, 55, 0.9], # L_Hip (11) [42, 55, 0.9], # R_Hip (12) [60, 75, 0.9], # L_Knee (13) [40, 75, 0.9], # R_Knee (14) [60, 90, 0.9], # L_Ankle (15) [40, 90, 0.9], # R_Ankle (16) ] pose_desc = edit_api._describe_pose(standing_kpts) self.assertIn("standing", pose_desc) self.assertIn("facing forward", pose_desc) gaze = edit_api._detect_facial_direction(standing_kpts) self.assertEqual(gaze, "looking forward") # Test 2: Gaze look left and up left_up_kpts = [list(kp) for kp in standing_kpts] # To look left: nose (0) moves to the left (smaller X than ear midpoint 50) left_up_kpts[0][0] = 45 # Nose X # To look up: nose (0) moves up (y-level close to eye level y=11) left_up_kpts[0][1] = 11.2 # Nose Y gaze_left_up = edit_api._detect_facial_direction(left_up_kpts) self.assertEqual(gaze_left_up, "looking left and up") # Test 3: Lying down / reclining pose lying_kpts = [list(kp) for kp in standing_kpts] # In lying down, nose is at y=10, hips are at y=12 (very horizontal) # Head X = 10, Hip X = 80 lying_kpts[0][0] = 10 # Head X lying_kpts[0][1] = 10 # Head Y lying_kpts[11][0] = 80 # Hip X lying_kpts[11][1] = 12 # Hip Y lying_kpts[12][0] = 80 lying_kpts[12][1] = 12 pose_desc_lying = edit_api._describe_pose(lying_kpts) self.assertIn("lying down", pose_desc_lying) def test_18_designer_enhancements(self): import unittest.mock as mock import edit_api # 1. Test Deduplication of Pose Names # We will mock the external requests.post LLM call to return some mock poses, # one of which collides with an existing pose name in poses.md (e.g., "The Clasp" or "standing") mock_raw_response = ( "# standing\n" "This is a standing pose description.\n" "Perfect anatomy, photorealistic.\n\n" "# Custom Pose\n" "This is a custom pose description.\n" "Anatomically precise, photorealistic.\n" ) mock_response = mock.Mock() mock_response.status_code = 200 mock_response.json.return_value = { "choices": [ { "message": { "content": mock_raw_response } } ] } with mock.patch("requests.post", return_value=mock_response) as mock_post: # Let's override _load_poses to return {"standing": "existing text"} to force a name collision with mock.patch("edit_api._load_poses", return_value={"standing": {"text": "existing text", "beta": False}}): payload = { "n": 2, "context": "standing and custom pose instructions", "filename": None, "beta": False, "messages": None } response = self.client.post("/designer/generate", json=payload) self.assertEqual(response.status_code, 200) data = response.json() self.assertEqual(data["status"], "success") poses = data["poses"] # The duplicate "standing" name should be renamed to "standing 2" instead of skipped! self.assertIn("standing 2", poses) self.assertIn("Custom Pose", poses) self.assertEqual(poses["standing 2"], "This is a standing pose description. Perfect anatomy, photorealistic.") self.assertEqual(poses["Custom Pose"], "This is a custom pose description. Anatomically precise, photorealistic.") # 2. Test Multi-Turn Conversation/History Payload Building with mock.patch("requests.post", return_value=mock_response) as mock_post: with mock.patch("edit_api._load_poses", return_value={}): # Send history messages in the request history_messages = [ {"role": "user", "content": "Initial prompt"}, {"role": "assistant", "content": "# Some Pose\nbody text"} ] payload_with_history = { "n": 1, "context": "make them sit down", "filename": None, "beta": False, "messages": history_messages } response = self.client.post("/designer/generate", json=payload_with_history) self.assertEqual(response.status_code, 200) # Check that requests.post was called with the correct messages payload self.assertTrue(mock_post.called) call_args = mock_post.call_args called_json = call_args[1]["json"] called_messages = called_json["messages"] # The payload should contain: # 1. System message (DESIGNER_SYSTEM) # 2. Initial user prompt # 3. Assistant response # 4. New user follow-up prompt incorporating "make them sit down" self.assertEqual(len(called_messages), 4) self.assertEqual(called_messages[0]["role"], "system") self.assertEqual(called_messages[1]["role"], "user") self.assertEqual(called_messages[1]["content"], "Initial prompt") self.assertEqual(called_messages[2]["role"], "assistant") self.assertEqual(called_messages[2]["content"], "# Some Pose\nbody text") self.assertEqual(called_messages[3]["role"], "user") self.assertIn("make them sit down", called_messages[3]["content"]) if __name__ == "__main__": unittest.main()