Files
Athena-Deck/test_music_components.py

48 lines
3.6 KiB
Python

import tempfile,unittest
from pathlib import Path
from unittest.mock import Mock
from profiles import Profiles,YUE2_REPO,YUE2_COMPONENTS,music_recipe
from catalog import file_role
class MusicComponentsTests(unittest.TestCase):
def test_download_recipe_assignment_and_remaining_worker_block(self):
model=dict(id='main',kind='music',repo=YUE2_REPO,file='yue2-3b-bf16.gguf',revision='a'*40)
entries=[dict(id=role,kind='music',repo=YUE2_REPO,file=info['files'][0],revision='a'*40,size=100) for role,info in YUE2_COMPONENTS.items()]
lookup={e['id']:e for e in [model]+entries}
c=Mock();c.entry.side_effect=lambda id:lookup[id] if id in lookup else (_ for _ in ()).throw(ValueError('missing'))
c.status.return_value=dict(entries=entries)
c.files.return_value=dict(repo=YUE2_REPO,revision='a'*40,files=[dict(name=e['file'],size=100) for e in entries],gated=False,license='apache-2.0')
with tempfile.TemporaryDirectory() as d:
p=Profiles(Path(d)/'profiles.json',c);p.music_runtime_ready=lambda:True
p.rows=[dict(id='profile',kind='music',revision=1,name='YuE2',model_id='main',parameters={})]
r=p.components('main');self.assertEqual(len(r['requirements']),5);self.assertTrue(r['runtime_installed']);c.files.assert_called_once_with(YUE2_REPO,'a'*40)
self.assertTrue(all(len(x['available'])==1 for x in r['requirements']))
p.assign(dict(id='profile',revision=1,components={e['id']:e['id'] for e in entries}))
result=p.status()['profiles'][0];self.assertFalse(result['runnable']);self.assertEqual(len(result['blockers']),1)
entries[0]['revision']='b'*40
with self.assertRaises(ValueError):p._component(model,'vae','vae')
def test_components_cannot_be_main_models(self):
for info in YUE2_COMPONENTS.values():
for filename in info['files']:self.assertFalse(file_role(filename)['profile_eligible'])
self.assertIsNone(music_recipe(dict(kind='music',repo='other/Yue2',file='yue2-3b-bf16.gguf')))
def test_manual_setup_available_in_every_category_and_preserves_recipe(self):
for kind in ('chat','image','audio','stt','music','voice','video'):
model=dict(id='main',kind=kind,repo='unknown/model',file='weights.gguf',revision='a'*40)
component=dict(id='aux',kind=kind,repo='external/repo',file='vae.safetensors',revision='b'*40,size=1)
c=Mock();c.entry.side_effect=lambda id: {'main':model,'aux':component}[id]
c.status.return_value=dict(entries=[component]);c.files.return_value=dict(repo='unknown/model',revision='a'*40,files=[],gated=False,url='https://huggingface.co/unknown/model',license=None)
with tempfile.TemporaryDirectory() as d:
p=Profiles(Path(d)/'profiles.json',c);p.rows=[dict(id='p',kind=kind,revision=1,name='Test',model_id='main',parameters={})]
self.assertTrue(p.components('main')['manual'])
p.assign(dict(id='p',revision=1,components={'manual_vae':'aux'}))
self.assertEqual(p.rows[0]['components']['manual_vae'],'aux')
self.assertTrue(p.references('aux'))
with self.assertRaises(ValueError):p.assign(dict(id='p',revision=2,components={'arbitrary':'aux'}))
def test_audio_runtime_package_definition_matches_exact_repo_and_model(self):
import json
from profiles import audio_package_recipe
with tempfile.TemporaryDirectory() as d:
Path(d,'model_specs').mkdir();Path(d,'model_specs/test.json').write_text(json.dumps(dict(packages=[dict(target_directory='Example-GGUF',files=['model.gguf','config.json'])])))
model=dict(repo='audio-cpp/Example-GGUF',file='model.gguf')
self.assertEqual(audio_package_recipe(model,d)['runtime_file_0']['files'],['config.json'])
self.assertIsNone(audio_package_recipe(dict(model,repo='other/Example-GGUF'),d))