57 lines
2.7 KiB
Go
57 lines
2.7 KiB
Go
package main
|
|
|
|
import (
|
|
"context"
|
|
"net/http"
|
|
"net/http/httptest"
|
|
"strings"
|
|
"testing"
|
|
|
|
"git.cin.su/m/Assistant/adapters/models"
|
|
"git.cin.su/m/Assistant/contracts/capability"
|
|
"git.cin.su/m/Assistant/contracts/cognitive"
|
|
)
|
|
|
|
type qualificationProvider struct{}
|
|
|
|
func (qualificationProvider) Probe(context.Context) models.ProviderProbe {
|
|
return models.ProviderProbe{Health: capability.HealthAvailable}
|
|
}
|
|
func (qualificationProvider) Discover(context.Context) ([]models.DiscoveredModel, models.ProviderProbe) {
|
|
return []models.DiscoveredModel{{ProviderModelID: "model-a"}, {ProviderModelID: "model-b"}}, models.ProviderProbe{Health: capability.HealthAvailable}
|
|
}
|
|
func (qualificationProvider) Advise(_ context.Context, request cognitive.AdvisoryRequest) (cognitive.AdvisoryResponse, error) {
|
|
return cognitive.AdvisoryResponse{Version: request.Version, RequestID: request.RequestID, AdvisoryText: "reachable", FinishReason: cognitive.FinishCompleted}, nil
|
|
}
|
|
|
|
func TestModelQualificationUsesDiscoveredModelsAndLeavesQualityUnrated(t *testing.T) {
|
|
provider := qualificationProvider{}
|
|
mux := http.NewServeMux()
|
|
mux.Handle("POST /v1/provider-discovery/{provider}/qualify", modelQualificationHandler(map[string]providerRegistration{"gpu": {ProviderInstanceID: "gpu", Provider: provider, Discovery: provider}}, "owner"))
|
|
server := httptest.NewServer(mux)
|
|
defer server.Close()
|
|
request := httptest.NewRequest(http.MethodPost, server.URL+"/v1/provider-discovery/gpu/qualify", strings.NewReader(`{"prompt":"ready","max_output_bytes":128}`))
|
|
request.Header.Set("Content-Type", "application/json")
|
|
response := httptest.NewRecorder()
|
|
mux.ServeHTTP(response, request)
|
|
if response.Code != http.StatusOK {
|
|
t.Fatalf("qualification status=%d body=%s", response.Code, response.Body.String())
|
|
}
|
|
if !strings.Contains(response.Body.String(), `"status":"REACHABLE"`) || !strings.Contains(response.Body.String(), `"category":"UNRATED"`) || strings.Contains(response.Body.String(), `"quality_assessed":true`) {
|
|
t.Fatalf("qualification overclaimed quality: %s", response.Body.String())
|
|
}
|
|
}
|
|
|
|
func TestModelQualificationRejectsUnboundedOutput(t *testing.T) {
|
|
provider := qualificationProvider{}
|
|
mux := http.NewServeMux()
|
|
mux.Handle("POST /v1/provider-discovery/{provider}/qualify", modelQualificationHandler(map[string]providerRegistration{"cpu": {ProviderInstanceID: "cpu", Provider: provider, Discovery: provider}}, "owner"))
|
|
request := httptest.NewRequest(http.MethodPost, "/v1/provider-discovery/cpu/qualify", strings.NewReader(`{"prompt":"ready","max_output_bytes":8192}`))
|
|
request.Header.Set("Content-Type", "application/json")
|
|
response := httptest.NewRecorder()
|
|
mux.ServeHTTP(response, request)
|
|
if response.Code != http.StatusBadRequest {
|
|
t.Fatalf("unbounded qualification status=%d", response.Code)
|
|
}
|
|
}
|