import json
from pathlib import Path
import tempfile
import unittest
from unittest.mock import patch,Mock
from image_test import ImageTests,workflow,select_gpus

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_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()