106 lines
5.5 KiB
Python
106 lines
5.5 KiB
Python
"""Offline authenticated model settings/discovery contracts."""
|
|
import json
|
|
import pytest
|
|
import httpx
|
|
from fastapi.testclient import TestClient
|
|
from control.app import Settings, create_app
|
|
from control.model_catalog import ModelCatalog, CatalogUnavailable
|
|
|
|
PASSWORD = 'correct-horse-password'
|
|
SECRET = 'internal-secret-' + 'x' * 32
|
|
|
|
class Provisioner:
|
|
async def ensure(self, user):
|
|
return {'container_name': user['container_name'], 'hermes_url': 'http://127.0.0.1:1234', 'hermes_api_key': 'private-hermes-key'}
|
|
async def ready(self, user): return True
|
|
async def remove(self, user): pass
|
|
|
|
class Catalog:
|
|
unavailable = False
|
|
async def models(self):
|
|
if self.unavailable: raise RuntimeError('secret-upstream-key')
|
|
return [{'id': 'auto', 'name': 'auto', 'thinking_levels': ['none', 'high']},
|
|
{'id': 'openai/test', 'name': 'openai/test', 'thinking_levels': ['none', 'high']}]
|
|
|
|
@pytest.fixture
|
|
def env(tmp_path):
|
|
catalog = Catalog()
|
|
app = create_app(Settings(db_path=str(tmp_path/'accounts.db'), admin_password=PASSWORD,
|
|
service_secret=SECRET, allowed_origins=('https://testserver',), reconcile_interval=3600), Provisioner(), catalog)
|
|
with TestClient(app, base_url='https://testserver') as client:
|
|
yield client, app, catalog
|
|
|
|
def login(client, name='admin'):
|
|
result = client.post('/api/auth/login', json={'username': name, 'password': PASSWORD}, headers={'Origin':'https://testserver'})
|
|
assert result.status_code == 200
|
|
return {'Authorization': 'Bearer ' + result.json()['token']}, result.json()['user']['id']
|
|
|
|
def test_authenticated_saved_settings_are_account_scoped(env):
|
|
client, app, catalog = env
|
|
assert client.get('/api/settings/model').status_code == 401
|
|
headers, uid = login(client)
|
|
assert client.get('/api/settings/model', headers=headers).json()['settings'] == {'model':'auto', 'thinking_level':'none'}
|
|
saved = {'model':'openai/test', 'thinking_level':'high'}
|
|
assert client.put('/api/settings/model', headers=headers, json=saved).json()['settings'] == saved
|
|
created = client.post('/api/admin/users', headers=headers, json={'username':'other', 'display_name':'Other', 'password':PASSWORD})
|
|
assert created.status_code == 201
|
|
other, _ = login(client, 'other')
|
|
assert client.get('/api/settings/model', headers=other).json()['settings'] == {'model':'auto', 'thinking_level':'none'}
|
|
assert client.get('/api/settings/model', headers=headers).json()['settings'] == saved
|
|
app.state.accounts.update(uid, status='ready', hermes_url='http://127.0.0.1:1234', hermes_api_key='private-hermes-key')
|
|
response = client.get('/internal/users/'+uid, headers={'X-Control-Secret':SECRET})
|
|
assert response.status_code == 200
|
|
assert response.json()['model_settings'] == saved
|
|
assert 'secret' not in json.dumps(client.get('/api/settings/model', headers=headers).json())
|
|
|
|
@pytest.mark.parametrize('body', [
|
|
{'model':'unlisted', 'thinking_level':'none'},
|
|
{'model':'auto', 'thinking_level':'ultra'},
|
|
{'model':'auto', 'thinking_level':'none', 'provider':'evil'},
|
|
{'model':'auto', 'thinking_level':'none', 'user_id':'another'},
|
|
{'model':'auto', 'thinking_level':{'enabled':False}},
|
|
])
|
|
def test_rejects_untrusted_model_effort_and_routing(env, body):
|
|
client, app, catalog = env
|
|
headers, uid = login(client)
|
|
assert client.put('/api/settings/model', headers=headers, json=body).status_code == 400
|
|
assert app.state.accounts.get_model_settings(uid) == {'model':'auto', 'thinking_level':'none'}
|
|
|
|
def test_csrf_and_catalog_failure_preserve_settings(env):
|
|
client, app, catalog = env
|
|
headers, uid = login(client)
|
|
saved = {'model':'openai/test', 'thinking_level':'high'}
|
|
assert client.put('/api/settings/model', headers=headers, json=saved).status_code == 200
|
|
assert client.put('/api/settings/model', json=saved).status_code == 403
|
|
assert client.put('/api/settings/model', headers={'Origin':'https://evil'}, json=saved).status_code == 403
|
|
catalog.unavailable = True
|
|
get = client.get('/api/settings/model', headers=headers)
|
|
assert get.status_code == 200 and not get.json()['catalog_available']
|
|
assert get.json()['settings'] == saved
|
|
put = client.put('/api/settings/model', headers=headers, json={'model':'auto','thinking_level':'none'})
|
|
assert put.status_code == 503 and 'secret' not in put.text
|
|
assert app.state.accounts.get_model_settings(uid) == saved
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_catalog_projects_metadata_and_caches():
|
|
calls = []
|
|
def handle(request):
|
|
calls.append(request)
|
|
return httpx.Response(200,json={'data':[{'id':'auto','api_key':'NEVER'}, {'id':'openai/test'}, {'id':'bad model'}, {'id':'auto'}]})
|
|
async with httpx.AsyncClient(transport=httpx.MockTransport(handle)) as client:
|
|
catalog = ModelCatalog(base_url='https://trusted/v1', api_key='SECRET', client=client)
|
|
first = await catalog.models()
|
|
assert [m['id'] for m in first] == ['auto','openai/test']
|
|
assert 'NEVER' not in json.dumps(first)
|
|
assert await catalog.models() == first and len(calls) == 1
|
|
assert calls[0].headers['Authorization'] == 'Bearer SECRET'
|
|
assert str(calls[0].url) == 'https://trusted/v1/models'
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_catalog_errors_are_redacted():
|
|
async with httpx.AsyncClient(transport=httpx.MockTransport(lambda r:httpx.Response(401,text='SECRET'))) as client:
|
|
catalog = ModelCatalog(api_key='SECRET', client=client)
|
|
with pytest.raises(CatalogUnavailable) as exc:
|
|
await catalog.models()
|
|
assert 'SECRET' not in str(exc.value)
|