import json from pathlib import Path import tempfile import unittest from unittest.mock import patch,Mock from image_test import ImageTests,workflow,select_gpus,decode_references class ImageTestTests(unittest.TestCase): def test_workflow_uses_selected_profile_and_owned_aliases(self): p=dict(width=512,height=768,steps=12,guidance=1.5) graph=workflow('benign test',p,99) self.assertEqual(graph['1']['inputs']['unet_name'],'model.gguf') self.assertEqual(graph['2']['class_type'],'DeckTextEncoderLoader') self.assertEqual(graph['6']['inputs']['seed'],99) self.assertEqual(graph['6']['inputs']['steps'],12) self.assertEqual(graph['6']['inputs']['cfg'],1.5) self.assertEqual(graph['5']['inputs']['height'],768) def test_qwen_references_feed_conditioning_and_edit_latent(self): p=dict(width=512,height=512,steps=4,guidance=1) graph=workflow('synthetic edit',p,42,references=['reference-1.png','reference-2.jpg']) self.assertEqual(graph['4']['inputs']['vae'],['3',0]) self.assertEqual(graph['4']['inputs']['images.image_1'],['21',0]) self.assertEqual(graph['4']['inputs']['images.image_2'],['22',0]) self.assertEqual(graph['6']['inputs']['latent_image'],['4',2]) self.assertEqual(graph['21']['inputs']['image'],'reference-1.png') self.assertNotIn('image_1',graph['4']['inputs']) with self.assertRaises(ValueError):workflow('test',p,1,'unknown',['reference-1.png']) def test_flux_references_feed_both_conditioning_branches(self): p=dict(width=512,height=512,steps=4,guidance=1) for count in (1,4): graph=workflow('synthetic edit',p,42,'flux',[f'reference-{i}.png' for i in range(count)]) self.assertEqual(sum(n['class_type']=='ReferenceLatent' for n in graph.values()),count*2) self.assertEqual(graph['10']['inputs']['positive'],[str(23+(count-1)*5),0]) self.assertEqual(graph['10']['inputs']['negative'],[str(24+(count-1)*5),0]) self.assertEqual(graph['22']['inputs']['vae'],['3',0]) self.assertEqual(graph['6']['inputs']['latent_image'],['5',0]) with self.assertRaises(ValueError):workflow('test',p,1,'flux',['r.png']*5) def test_reference_capabilities_follow_supported_recipes(self): from profiles import image_capabilities,FLUX_REPO,QWEN_REPO self.assertEqual(image_capabilities(dict(repo=FLUX_REPO,file='Flux.2 Klein-9B_fp16_nsfw.safetensors'))['reference_images'],4) self.assertEqual(image_capabilities(dict(repo=QWEN_REPO,file='model.gguf'))['reference_images'],4) self.assertEqual(image_capabilities(dict(repo='unknown',file='model.safetensors'))['reference_images'],0) def test_family_registry_reuses_workflow_from_lineage_not_filename(self): from profiles import image_adapter,image_recipe,FLUX_COMPONENTS m=dict(repo='someone/new-weights',file='custom.safetensors',runtime_metadata={'base_models':['black-forest-labs/FLUX.2-klein-9B']}) self.assertEqual(image_adapter(m)['workflow'],'flux') self.assertEqual(image_recipe(m),FLUX_COMPONENTS) m['runtime_metadata']={};m['file']='flux-qwen-looking-name.safetensors' self.assertIsNone(image_adapter(m)) m['runtime_metadata']={'architecture':'Flux2Transformer2DModel'} self.assertIsNone(image_adapter(m)) m['runtime_metadata']={'architecture':'qwen_image_2_1'};m['file']='new.gguf' self.assertEqual(image_adapter(m)['workflow'],'qwen') m['file']='vae/new.gguf';self.assertIsNone(image_adapter(m)) def test_reference_validation_enforces_profile_limit_and_signature(self): import base64 png='data:image/png;base64,'+base64.b64encode(b'\x89PNG\r\n\x1a\nsynthetic').decode() self.assertEqual(len(decode_references([png],4)),1) for value,limit in [([png],0),([png]*5,4),(['data:image/png;base64,'+base64.b64encode(b'not png').decode()],4)]: with self.assertRaises(ValueError):decode_references(value,limit) def test_flux_graph_and_offload(self): p=dict(width=512,height=512,steps=4,guidance=1) graph=workflow('a red bicycle',p,42,'flux') self.assertEqual(graph['1']['class_type'],'UNETLoader');self.assertEqual(graph['2']['inputs']['family'],'flux2') self.assertEqual(graph['12']['class_type'],'Flux2Scheduler');self.assertEqual(graph['12']['inputs']['steps'],4) self.assertEqual(graph['5']['class_type'],'EmptyFlux2LatentImage') devices=[dict(name='RTX 5080',uuid='a',free_mib=16000,processes=0),dict(name='RTX 3060',uuid='b',free_mib=12000,processes=0)] self.assertEqual(select_gpus(devices,18*1024**3,8*1024**3,400*1024**2,offload=True)[0]['uuid'],'a') with self.assertRaises(ValueError):select_gpus(devices,18*1024**3,8*1024**3,400*1024**2) def test_two_gpu_selection_rejects_busy_or_small_encoder_gpu(self): devices=[dict(name='NVIDIA GeForce RTX 3060',uuid='encoder',processes=0,free_mib=11962),dict(name='NVIDIA GeForce RTX 5080',uuid='image',processes=0,free_mib=15923)] image,encoder=select_gpus(devices,7*1024**3,9*1024**3,700*1024**2) self.assertEqual((image['uuid'],encoder['uuid']),('image','encoder')) for change in (dict(processes=1),dict(free_mib=7000)): with self.assertRaises(ValueError):select_gpus([dict(devices[0],**change),devices[1]],7*1024**3,9*1024**3,700*1024**2) with self.assertRaises(ValueError):select_gpus([devices[0],dict(devices[1],processes=1)],7*1024**3,9*1024**3,700*1024**2) def test_missing_runtime_and_invalid_prompts_do_not_spawn(self): with tempfile.TemporaryDirectory() as d,patch('image_test.PYTHON',Path(d)/'missing'),patch('image_test.subprocess.Popen') as spawn: manager=ImageTests(d,Mock()) for prompt in ('',None,'x'*4001,'hello'): with self.assertRaises(ValueError):manager.start('profile',prompt) spawn.assert_not_called() def test_images_only_accessible_for_current_completed_job(self): with tempfile.TemporaryDirectory() as d: manager=ImageTests(d,Mock());manager.job=dict(id='a'*32,state='complete') p=Path(d)/('a'*32);p.mkdir();(p/'result.png').write_bytes(b'png') self.assertEqual(manager.image('a'*32),b'png') for value in ('../../secret','b'*32): with self.assertRaises(ValueError):manager.image(value) manager.job['state']='running' with self.assertRaises(ValueError):manager.image('a'*32) def test_restart_marks_active_job_interrupted(self): with tempfile.TemporaryDirectory() as d: (Path(d)/'status.json').write_text(json.dumps(dict(id='test',state='running'))) manager=ImageTests(d,Mock());self.assertEqual(manager.status()['job']['state'],'interrupted') def test_cancel_terminates_only_owned_process(self): import subprocess,sys with tempfile.TemporaryDirectory() as d: manager=ImageTests(d,Mock()) child=subprocess.Popen([sys.executable,'-c','import time;time.sleep(60)'],start_new_session=True) manager.process=child try: self.assertTrue(manager.stop()['cancellation_requested']);self.assertIsNotNone(child.poll()) finally: if child.poll() is None:child.kill();child.wait() def test_busy_job_does_not_start_second_process(self): with tempfile.TemporaryDirectory() as d,patch('image_test.subprocess.Popen') as spawn: manager=ImageTests(d,Mock());manager.job={'state':'running'} with self.assertRaises(ValueError):manager.start('profile','hello') spawn.assert_not_called()