54 lines
3.5 KiB
Python
54 lines
3.5 KiB
Python
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'])
|