49 lines
2.6 KiB
Python
49 lines
2.6 KiB
Python
|
import asyncio
|
||
|
|
import json
|
||
|
|
from types import SimpleNamespace
|
||
|
|
import httpx
|
||
|
|
import pytest
|
||
|
|
from fastapi import HTTPException, BackgroundTasks
|
||
|
|
from app.settings.backend import ai_benchmark as b
|
||
|
|
from app.settings.backend import router as r
|
||
|
|
|
||
|
|
@pytest.mark.parametrize('url',['file:///tmp/a','http://user:secret@host/v1','https://host/v1?key=secret','https://host/#key','host'])
|
||
|
|
def test_bad_endpoint_rejected(url):
|
||
|
|
with pytest.raises(HTTPException): b.validate_endpoint(url)
|
||
|
|
|
||
|
|
def test_endpoint_keeps_local_url():
|
||
|
|
assert b.validate_endpoint(' http://172.16.31.195:8080/v1/ ') == 'http://172.16.31.195:8080/v1'
|
||
|
|
|
||
|
|
@pytest.mark.parametrize('provider',['openai','ollama'])
|
||
|
|
def test_both_providers_use_same_prompts_and_checks(monkeypatch,provider):
|
||
|
|
calls=[]
|
||
|
|
class Client:
|
||
|
|
def __init__(self,**kw): pass
|
||
|
|
async def __aenter__(self): return self
|
||
|
|
async def __aexit__(self,*args): pass
|
||
|
|
async def post(self,url,**kw):
|
||
|
|
calls.append((url,kw))
|
||
|
|
envelope={'choices':[{'message':{'content':'{"answer":42}'}}],'usage':{'prompt_tokens':5,'completion_tokens':7}} if provider=='openai' else {'message':{'content':'{"answer":42}'},'prompt_eval_count':5,'eval_count':7}
|
||
|
|
return httpx.Response(200,json=envelope)
|
||
|
|
monkeypatch.setattr(b.httpx,'AsyncClient',Client)
|
||
|
|
test={'prompt':'Same test','checks':[('Answer',lambda x:x['answer']==42)]}
|
||
|
|
result=asyncio.run(b._run_case('http://host/v1','model',test,provider,'dummy-secret'))
|
||
|
|
assert result['score']==1 and result['response_tokens']==7
|
||
|
|
assert calls[0][1]['json']['messages'][1]['content']=='Same test'
|
||
|
|
assert calls[0][1]['headers']['Authorization']=='Bearer dummy-secret'
|
||
|
|
assert calls[0][0].endswith('/chat/completions' if provider=='openai' else '/api/chat')
|
||
|
|
|
||
|
|
def test_alternative_only_does_not_need_current_server_and_never_stores_key(monkeypatch):
|
||
|
|
captured=[]
|
||
|
|
monkeypatch.setattr(r,'execute_query_single',lambda *args:None)
|
||
|
|
monkeypatch.setattr(r,'execute_query',lambda sql,args:captured.append(args) or [{'id':1}])
|
||
|
|
async def no_fetch(*args): pytest.fail('current server must not be needed')
|
||
|
|
monkeypatch.setattr(b,'fetch_models',no_fetch)
|
||
|
|
tasks=BackgroundTasks()
|
||
|
|
payload=r.AIBenchmarkRunRequest(models=[],alternative_url='http://host/v1',alternative_model='model',alternative_api_key='dummy-secret')
|
||
|
|
asyncio.run(r.create_ai_benchmark_run(payload,tasks,{'id':1}))
|
||
|
|
assert 'dummy-secret' not in str(captured)
|
||
|
|
assert 'dummy-secret' not in repr(payload)
|
||
|
|
assert len(json.loads(captured[0][0]))==1
|
||
|
|
assert tasks.tasks[0].args[-1]['api_key']=='dummy-secret'
|