import tempfile,time,unittest
from pathlib import Path
from unittest.mock import Mock,patch
from music import Music
from profiles import YUE2_COMPONENTS,YUE2_REPO
class MusicTests(unittest.TestCase):
 def make(self,root):
  root=Path(root);catalog=Mock();catalog.root=root/'models';catalog.root.mkdir()
  main=dict(id='main',repo=YUE2_REPO,file='yue2-3b-bf16.gguf',kind='music',size=1)
  entries={role:dict(id=role,file=info['files'][0]) for role,info in YUE2_COMPONENTS.items()}
  for entry in [main]+list(entries.values()):
   (catalog.root/entry['id']).mkdir();(catalog.root/entry['id']/('model'+Path(entry['file']).suffix)).touch()
  catalog.entry.side_effect=lambda id:entries[id]
  p=dict(id='p',kind='music',runnable=True,model=main,components={role:role for role in entries})
  profiles=Mock();profiles.status.return_value=dict(profiles=[p]);profiles._component.side_effect=lambda model,role,id:entries[id]
  cli=root/'cli';cli.write_text('#!/usr/bin/env python3\nimport sys,wave\np=sys.argv[sys.argv.index("--out")+1]\nwith wave.open(p,"wb") as w:w.setparams((1,2,8000,0,"NONE","none"));w.writeframes(bytes(1600))\n');cli.chmod(0o755)
  runtime=Mock();runtime.status.return_value=dict(installed=True);runtime.paths.return_value=(cli,root)
  scheduler=Mock();release=Mock();scheduler.image_reservation.return_value=release
  return Music(root/'jobs',profiles,catalog,runtime,scheduler),release,p
 def wait(self,m):
  for _ in range(80):
   if m.status()['job']['state']!='running':return
   time.sleep(.05)
  self.fail('Test timeout')
 def test_success_and_output_without_retaining_prompts(self):
  with tempfile.TemporaryDirectory() as root:
   m,release,p=self.make(root)
   with patch('music.probe',return_value=[dict(name='NVIDIA RTX 5080',uuid='GPU-test',free_mib=16000,processes=0)]):
    job=m.start('p','Synthetic lyrics','Synthetic pop',max_tokens=32,steps=1);self.wait(m)
   self.assertEqual(m.status()['job']['state'],'complete');self.assertTrue(m.audio(job['id']).startswith(b'RIFF'));release.assert_called_once()
   self.assertFalse((m.root/job['id']/'worker.tmp').exists());self.assertNotIn('lyrics',m.status()['job'])
 def test_occupied_foreign_gpu_releases_lease(self):
  with tempfile.TemporaryDirectory() as root:
   m,release,p=self.make(root)
   with patch('music.probe',return_value=[dict(name='NVIDIA RTX 5080',uuid='GPU-test',free_mib=16000,processes=1)]):m.start('p','Synthetic','pop');self.wait(m)
   self.assertEqual(m.status()['job']['state'],'failed');self.assertIn('Fremde',m.status()['job']['phase']);release.assert_called_once()
 def test_invalid_request_and_missing_components(self):
  with tempfile.TemporaryDirectory() as root:
   m,release,p=self.make(root)
   with self.assertRaises(ValueError):m.start('p','Synthetic','pop',max_tokens=True)
   m.profiles._component.side_effect=ValueError('missing');self.assertEqual(len(m.blockers(p)),5)
   with self.assertRaises(ValueError):m.audio('../outside')

 def test_cancel_terminates_owned_process_and_releases_gpu(self):
  with tempfile.TemporaryDirectory() as root:
   m,release,p=self.make(root);m.runtime.paths()[0].write_text('#!/usr/bin/env python3\nimport time;time.sleep(30)\n')
   with patch('music.probe',return_value=[dict(name='NVIDIA RTX 5080',uuid='GPU-test',free_mib=16000,processes=0)]):
    m.start('p','Synthetic','pop')
    for _ in range(50):
     if m.process:break
     time.sleep(.02)
    m.stop();self.wait(m)
   self.assertEqual(m.status()['job']['state'],'cancelled');release.assert_called_once();self.assertFalse(m.status()['loaded'])