123 lines
8.2 KiB
Python
123 lines
8.2 KiB
Python
import json
|
|
import threading
|
|
import unittest
|
|
import tempfile
|
|
import secrets
|
|
import time
|
|
import urllib.request
|
|
import urllib.error
|
|
from unittest.mock import patch
|
|
from server import Server, HardwareProvider
|
|
|
|
class Tests(unittest.TestCase):
|
|
def setUp(self):
|
|
self.state=tempfile.TemporaryDirectory()
|
|
self.server=Server(0,state_dir=self.state.name)
|
|
self.password=secrets.token_urlsafe(32)
|
|
self.api_token=secrets.token_urlsafe(32)
|
|
record=self.server.credentials.setup(self.password,self.api_token)
|
|
self.server.sessions['test-session']=dict(expires=time.monotonic()+300,revision=record['password']['hash'])
|
|
self.cookie='deck_session=test-session'
|
|
self.thread=threading.Thread(target=self.server.serve_forever)
|
|
self.thread.start()
|
|
self.url=f'http://127.0.0.1:{self.server.server_port}'
|
|
def tearDown(self):
|
|
self.server.shutdown()
|
|
self.server.endpoint.close()
|
|
self.server.server_close()
|
|
self.thread.join()
|
|
self.state.cleanup()
|
|
def request(self,path,method='GET',headers=None):
|
|
req=urllib.request.Request(self.url+'/api/v1/'+path,method=method,headers={"Cookie":self.cookie,**(headers or {})})
|
|
with urllib.request.urlopen(req) as r:return json.load(r)
|
|
def test_demo_removed_from_status(self):
|
|
state=self.request('status')
|
|
self.assertNotIn('demo',state)
|
|
self.assertEqual(state['endpoint']['state'],'stopped')
|
|
with self.assertRaises(urllib.error.HTTPError) as exc:self.request('demo')
|
|
self.assertEqual(exc.exception.code,404);exc.exception.close()
|
|
def test_profile_and_download_routes(self):
|
|
from pathlib import Path
|
|
target=Path(self.state.name)/'models'/('a'*64);target.mkdir(parents=True)
|
|
(target/'model.gguf').write_bytes(b'GGUFtest')
|
|
(target/'entry.json').write_text(json.dumps(dict(repo='test/image',file='model.gguf',size=8,kind='image',revision='b'*40)))
|
|
payload=dict(id=None,revision=0,name='image-test',kind='image',model_id='a'*64,parameters=dict(width=512,height=512,steps=10,seed=-1,guidance=1))
|
|
req=urllib.request.Request(self.url+'/api/v1/profiles/save',data=json.dumps(payload).encode(),headers={'Cookie':self.cookie,'X-Athena-Deck':'1','Content-Type':'application/json'})
|
|
with urllib.request.urlopen(req) as r:self.assertEqual(r.status,200)
|
|
row=self.request('profiles')['profiles'][0]
|
|
self.assertEqual(row['name'],'image-test')
|
|
delete=urllib.request.Request(self.url+'/api/v1/profiles/delete',data=json.dumps(dict(id=row['id'],revision=row['revision'])).encode(),headers={'Cookie':self.cookie,'X-Athena-Deck':'1','Content-Type':'application/json'})
|
|
self.server.image_tests.job=dict(state='running',profile_id=row['id'])
|
|
with self.assertRaises(urllib.error.HTTPError) as busy:urllib.request.urlopen(delete)
|
|
self.assertEqual(busy.exception.code,400);busy.exception.close()
|
|
self.server.image_tests.job=None
|
|
with urllib.request.urlopen(delete) as r:self.assertEqual(r.status,200)
|
|
self.assertEqual(self.request('profiles')['profiles'],[])
|
|
self.assertTrue((target/'model.gguf').exists())
|
|
with self.assertRaises(urllib.error.HTTPError) as error:
|
|
urllib.request.urlopen(urllib.request.Request(self.url+'/api/v1/profiles',headers={'Authorization':'Bearer '+self.api_token}))
|
|
self.assertEqual(error.exception.code,401);error.exception.close()
|
|
|
|
def test_image_test_routes_are_admin_only(self):
|
|
self.assertIn('installed',self.request('image-runtime'))
|
|
self.assertIn('runtime_installed',self.request('image-tests'))
|
|
with self.assertRaises(urllib.error.HTTPError) as error:
|
|
urllib.request.urlopen(urllib.request.Request(self.url+'/api/v1/image-tests',headers={'Authorization':'Bearer '+self.api_token}))
|
|
self.assertEqual(error.exception.code,401);error.exception.close()
|
|
with self.assertRaises(urllib.error.HTTPError) as error:self.request('image-tests/image?id=../../auth.json')
|
|
self.assertEqual(error.exception.code,404);error.exception.close()
|
|
|
|
def test_docker_admin_auth_and_explicit_install_confirmation(self):
|
|
with patch.object(self.server.docker,'status',return_value={'installed':True,'services':[]}):
|
|
self.assertTrue(self.request('docker')['installed'])
|
|
with self.assertRaises(urllib.error.HTTPError) as error:
|
|
urllib.request.urlopen(urllib.request.Request(self.url+'/api/v1/docker',headers={'Authorization':'Bearer '+self.api_token}))
|
|
self.assertEqual(error.exception.code,401);error.exception.close()
|
|
with patch.object(self.server.docker,'install',return_value={'started':True}) as install:
|
|
for payload,expected in [({},400),({'confirm':True},200)]:
|
|
req=urllib.request.Request(self.url+'/api/v1/docker/install',data=json.dumps(payload).encode(),headers={'Cookie':self.cookie,'X-Athena-Deck':'1','Content-Type':'application/json'})
|
|
try:
|
|
with urllib.request.urlopen(req) as result:self.assertEqual(result.status,expected)
|
|
except urllib.error.HTTPError as error:self.assertEqual(error.code,expected);error.close()
|
|
install.assert_called_once()
|
|
def test_endpoint_admin_only_configuration(self):
|
|
self.assertEqual(self.request('endpoint')['state'],'stopped')
|
|
url=self.url+'/api/v1/endpoint/config'
|
|
for headers,expected in [({'Authorization':'Bearer '+self.api_token,'X-Athena-Deck':'1'},401),({'Cookie':self.cookie},403)]:
|
|
req=urllib.request.Request(url,data=b'{"port":8120}',headers={'Content-Type':'application/json',**headers})
|
|
with self.assertRaises(urllib.error.HTTPError) as exc:urllib.request.urlopen(req)
|
|
self.assertEqual(exc.exception.code,expected);exc.exception.close()
|
|
with patch.object(self.server.endpoint,'configure',return_value={'port':8120}) as configure:
|
|
req=urllib.request.Request(url,data=b'{"port":8120}',headers={'Cookie':self.cookie,'X-Athena-Deck':'1','Content-Type':'application/json'})
|
|
with urllib.request.urlopen(req) as r:self.assertEqual(json.load(r)['port'],8120)
|
|
configure.assert_called_once_with({'port':8120})
|
|
|
|
def test_chat_test_routes_are_admin_only(self):
|
|
self.assertIsNone(self.request('chat-tests')['job'])
|
|
for path,method,body in [('chat-tests','GET',None),('chat-tests/start','POST',b'{}'),('auto-tests','GET',None),('auto-tests/start','POST',b'{}'),('auto-tests/save','POST',b'{}')]:
|
|
req=urllib.request.Request(self.url+'/api/v1/'+path,method=method,data=body,headers={'Authorization':'Bearer '+self.api_token,'X-Athena-Deck':'1','Content-Type':'application/json'})
|
|
with self.assertRaises(urllib.error.HTTPError) as exc:urllib.request.urlopen(req)
|
|
self.assertEqual(exc.exception.code,401);exc.exception.close()
|
|
with patch.object(self.server.chat_tests,'start',return_value={'state':'running'}) as start:
|
|
data={'profile_id':'test','messages':[{'role':'user','content':'synthetic'}],'max_tokens':8}
|
|
req=urllib.request.Request(self.url+'/api/v1/chat-tests/start',data=json.dumps(data).encode(),headers={'Cookie':self.cookie,'X-Athena-Deck':'1','Content-Type':'application/json'})
|
|
with urllib.request.urlopen(req) as r:self.assertEqual(json.load(r)['state'],'running')
|
|
start.assert_called_once_with(data)
|
|
|
|
def test_control_guard(self):
|
|
for headers in ({},{'X-Athena-Deck':'1','Origin':'http://evil.invalid'}):
|
|
with self.assertRaises(urllib.error.HTTPError) as e:self.request('demo/start','POST',headers)
|
|
self.assertEqual(e.exception.code,403)
|
|
e.exception.close()
|
|
with self.assertRaises(urllib.error.HTTPError) as e:self.request('models/start','POST',{'X-Athena-Deck':'1'})
|
|
self.assertEqual(e.exception.code,404)
|
|
e.exception.close()
|
|
def test_unavailable_hardware(self):
|
|
with patch.dict('os.environ', {'DECK_LOCAL_HARDWARE':'0'}), patch('server.subprocess.run',side_effect=OSError()):
|
|
result=HardwareProvider().snapshot()
|
|
self.assertFalse(result['available'])
|
|
self.assertEqual(result['gpus'],[])
|
|
self.assertIsNone(result['sampled_at'])
|
|
|
|
if __name__=='__main__':unittest.main()
|