103 lines
9.4 KiB
Python
103 lines
9.4 KiB
Python
import copy,hashlib,json,os,tempfile,time,unittest
|
|
from pathlib import Path
|
|
from unittest.mock import Mock,patch
|
|
from backup import Backup,validate
|
|
from backup_codec import seal,open_backup,MAX
|
|
from deploy.docker_backup import validate as validate_container,restore as restore_container
|
|
from server import Server
|
|
from profiles import SCHEMAS,CHAT_GPU_DEFAULTS,CHAT_SAMPLING
|
|
|
|
class CodecTests(unittest.TestCase):
|
|
def test_authenticated_roundtrip_wrong_password_and_tampering(self):
|
|
doc={'secret':'private-key-and-hf-token','settings':{'a':1}}
|
|
raw=seal(doc,'sixteen character secret')
|
|
self.assertNotIn(b'private-key',raw);self.assertEqual(open_backup(raw,'sixteen character secret'),doc)
|
|
for broken,key in [(raw,'another secure password'),(raw[:-1]+bytes([raw[-1]^1]),'sixteen character secret')]:
|
|
with self.assertRaisesRegex(ValueError,'falsch|beschädigt'):open_backup(broken,key)
|
|
self.assertNotEqual(seal(doc,'sixteen character secret'),raw)
|
|
def test_unbounded_compression_rejected(self):
|
|
import zlib,secrets
|
|
from cryptography.hazmat.primitives.ciphers.aead import AESGCM
|
|
from backup_codec import MAGIC
|
|
salt=secrets.token_bytes(16);nonce=secrets.token_bytes(12);head=MAGIC+salt+nonce
|
|
key=hashlib.scrypt(b'sixteen character secret',salt=salt,n=32768,r=8,p=1,maxmem=64*1024**2,dklen=32)
|
|
raw=head+AESGCM(key).encrypt(nonce,zlib.compress(b' '* (MAX+1)),head)
|
|
with self.assertRaises(ValueError):open_backup(raw,'sixteen character secret')
|
|
|
|
class RestoreTests(unittest.TestCase):
|
|
def setUp(self):
|
|
self.env=patch.dict(os.environ,{'DECK_NETWORK_MODE':'disabled'});self.env.start();self.tmp=tempfile.TemporaryDirectory();self.root=Path(self.tmp.name)
|
|
self.server=Server(0,state_dir=self.root);self.server.credentials.setup('test administrator password','test_api_token_12345678901234567890')
|
|
self.server.docker=Mock();self.server.docker.call.return_value={'services':[],'configured':True};self.server.docker.status.return_value={'services':[]}
|
|
self.disk=patch('backup.shutil.disk_usage',return_value=Mock(free=1024**4));self.disk.start()
|
|
self.server.runtime.prerequisites=Mock(return_value={'cuda_ready':True,'cpu_ready':True})
|
|
def tearDown(self):
|
|
self.server.backup.stop();self.server.endpoint.close();self.server.video.close();self.server.dashboard.close();self.server.server_close();self.tmp.cleanup();self.env.stop();self.disk.stop()
|
|
def model(self):
|
|
data=b'model-weights-not-in-backup';repo='test/model';revision='a'*40;filename='model.gguf';ident=hashlib.sha256((repo+revision+filename).encode()).hexdigest();p=self.root/'models'/ident;p.mkdir(parents=True)
|
|
(p/'model.gguf').write_bytes(data);entry=dict(repo=repo,revision=revision,file=filename,size=len(data),sha256=hashlib.sha256(data).hexdigest(),kind='chat',state='downloaded');(p/'entry.json').write_text(json.dumps(entry));return ident
|
|
def profile(self,ident):
|
|
params={k:v[2] for k,v in SCHEMAS['chat'].items()};params.update(CHAT_GPU_DEFAULTS);params.update(CHAT_SAMPLING)
|
|
self.server.profiles.save(dict(id=None,revision=0,name='Test Model',kind='chat',model_id=ident,parameters=params))
|
|
def wait(self):
|
|
for _ in range(100):
|
|
if not self.server.backup.busy():return self.server.backup.status()['job']
|
|
time.sleep(.05)
|
|
self.fail('restore did not finish')
|
|
def test_export_excludes_weights_media_logs_and_includes_settings_and_history(self):
|
|
ident=self.model();self.profile(ident)
|
|
for folder in ('image-tests','tts-tests','llama-worker'):
|
|
p=self.root/folder;p.mkdir(exist_ok=True);(p/'private.log').write_text('PRIVATE-CHAT-LOG')
|
|
doc=open_backup(self.server.backup.export('backup secure password'),'backup secure password')
|
|
self.assertNotIn('PRIVATE-CHAT-LOG',json.dumps(doc));self.assertNotIn('model-weights-not-in-backup',json.dumps(doc));self.assertIn('history',doc);validate(doc)
|
|
self.assertEqual(doc['models'][0]['id'],ident);self.assertEqual(len(doc['settings']['profiles.json']),1)
|
|
def test_restore_reuses_verified_models_restores_profiles_and_keeps_login_when_selected(self):
|
|
ident=self.model();self.profile(ident);raw=self.server.backup.export('backup secure password');self.server.profiles.rows=[];(self.root/'profiles.json').write_text('[]')
|
|
current=self.server.credentials.read();summary=self.server.backup.inspect(raw,'backup secure password');self.assertFalse(summary['blockers'])
|
|
self.server.backup.start(dict(id=summary['id'],services=[],confirm=True,restore_credentials=False));job=self.wait()
|
|
self.assertEqual(job['state'],'complete',job);self.assertEqual(self.server.profiles.rows[0]['name'],'Test Model');self.assertEqual(self.server.credentials.read(),current);self.assertTrue(any(x['state']=='reused' for x in job['items']));self.assertEqual(len(list((self.root/'recovery').glob('before-*.adbackup'))),1)
|
|
def test_empty_target_downloads_pinned_model_then_restores(self):
|
|
ident=self.model();self.profile(ident);doc=self.server.backup.snapshot();data=(self.root/'models'/ident/'model.gguf').read_bytes()
|
|
import shutil
|
|
shutil.rmtree(self.root/'models'/ident);self.server.profiles.rows=[];(self.root/'profiles.json').write_text('[]')
|
|
def download(repo,filename,revision,kind):
|
|
self.assertEqual((repo,filename,revision),('test/model','model.gguf','a'*40));p=self.root/'models'/ident;p.mkdir();(p/'model.gguf').write_bytes(data);m=dict(doc['models'][0]);m.pop('id');(p/'entry.json').write_text(json.dumps(m));self.server.catalog.job={'state':'complete'}
|
|
with patch.object(self.server.catalog,'start',side_effect=download):
|
|
summary=self.server.backup.inspect(seal(doc,'backup secure password'),'backup secure password');self.server.backup.start(dict(id=summary['id'],services=[],confirm=True,restore_credentials=True));job=self.wait()
|
|
self.assertEqual(job['state'],'complete',job);self.assertEqual(self.server.profiles.rows[0]['model_id'],ident)
|
|
def test_rejects_unsafe_paths_and_bad_references_before_changes(self):
|
|
ident=self.model();self.profile(ident);doc=self.server.backup.snapshot(include_history=False)
|
|
for change in (lambda d:d['models'][0].update(file='../secret'),lambda d:d['settings']['profiles.json'][0].update(model_id='f'*64),lambda d:d.update(version=99)):
|
|
bad=copy.deepcopy(doc);change(bad)
|
|
with self.assertRaises(ValueError):self.server.backup.inspect(seal(bad,'backup secure password'),'backup secure password')
|
|
self.assertEqual(len(self.server.profiles.rows),1)
|
|
def test_wrong_download_hash_marks_failure_without_claiming_complete(self):
|
|
ident=self.model();doc=self.server.backup.snapshot(include_history=False);(self.root/'models'/ident/'model.gguf').write_bytes(b'x'*len(b'model-weights-not-in-backup'))
|
|
summary=self.server.backup.inspect(seal(doc,'backup secure password'),'backup secure password');self.server.backup.start(dict(id=summary['id'],services=[],confirm=True,restore_credentials=False));job=self.wait();self.assertEqual(job['state'],'failed');self.assertNotEqual(job['phase'],'Wiederherstellung abgeschlossen. Oberfläche ist bereit.')
|
|
|
|
class ContainerTests(unittest.TestCase):
|
|
def config(self):return dict(name='test-ui',image='example/ui:1',registry='registry.example/ui@sha256:'+'a'*64,labels={'io.athena-deck.managed':'true','io.athena-deck.role':'application'},unsupported=False,network='bridge',ports={'8000/tcp':[{'HostIp':'127.0.0.1','HostPort':'8128'}]},restart='unless-stopped',env=[],entrypoint=None,cmd=None,user='65534',workdir='',memory=0,nanocpus=0,mounts=[],running=True)
|
|
def test_rejects_privileged_host_socket_public_ports_and_unknown_host_network(self):
|
|
c=self.config();validate_container(c)
|
|
for change in (lambda x:x.update(unsupported=True),lambda x:x.update(network='host'),lambda x:x['ports']['8000/tcp'][0].update(HostIp='0.0.0.0'),lambda x:x.update(mounts=[dict(index=0,target='/var/run/docker.sock',deck_path=None,read_only=True,was_file=True,files={})])):
|
|
bad=copy.deepcopy(c);change(bad)
|
|
with self.assertRaises(ValueError):validate_container(bad)
|
|
def test_only_deskweb_can_mount_its_specific_writable_media_folder(self):
|
|
c=self.config();c.update(name='ltx-deskweb',mounts=[dict(index=0,target='/data/inputs',deck_path='video/original-work/remote-inputs/deskweb',read_only=False,was_file=False,files={})]);validate_container(c)
|
|
for change in (lambda x:x.update(name='other-ui'),lambda x:x['mounts'][0].update(deck_path='models'),lambda x:x['mounts'][0].update(deck_path='video/original-work'),lambda x:x['mounts'][0].update(target='/data/outputs')):
|
|
bad=copy.deepcopy(c);change(bad)
|
|
with self.assertRaises(ValueError):validate_container(bad)
|
|
def test_ladypoly_recipe_and_config_are_scoped(self):
|
|
c=self.config();c.update(name='athena-ladypoly',image='athena-ladypoly:3960afc-deck1',registry=None,network='host',build_recipe='ladypoly',mounts=[dict(index=0,target='/data',deck_path=None,read_only=False,was_file=False,files={'deck-connection.json':'e30='})]);validate_container(c)
|
|
for change in (lambda x:x.update(name='other-ui'),lambda x:x.update(image='untrusted:1'),lambda x:x['mounts'][0].update(target='/other')):
|
|
bad=copy.deepcopy(c);change(bad)
|
|
with self.assertRaises(ValueError):validate_container(bad)
|
|
def test_never_replaces_existing_foreign_container(self):
|
|
with tempfile.TemporaryDirectory() as d:
|
|
manager=Mock(state=Path(d));(Path(d)/'backup-policy.json').write_text(json.dumps({'deck_state':d}))
|
|
x={'Config':{'Labels':{},'Image':'example/ui:1'}}
|
|
with patch('deploy.docker_backup.os.geteuid',return_value=0),patch('deploy.docker_backup.subprocess.run',return_value=Mock(returncode=0,stdout=json.dumps([x]))):
|
|
with self.assertRaises(ValueError):restore_container(manager,self.config())
|
|
|
|
if __name__=='__main__':unittest.main()
|