52 lines
2.5 KiB
Python
52 lines
2.5 KiB
Python
import json
|
|
from pathlib import Path
|
|
import tempfile
|
|
import unittest
|
|
from unittest.mock import patch,Mock
|
|
from image_test import ImageTests,workflow
|
|
|
|
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']['inputs']['device'],'cpu')
|
|
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_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()
|