quorate/services/agentgateway-model-catalog-sync/main_test.go

630 lines
19 KiB
Go
Raw Normal View History

package main
import (
"context"
"encoding/json"
"errors"
"io"
"net/http"
"net/http/httptest"
"os"
"path/filepath"
"reflect"
"slices"
"strings"
"testing"
"time"
)
func TestDurationFromEnvironment(t *testing.T) {
t.Setenv("TEST_INTERVAL", "")
got, err := durationFromEnvironment("TEST_INTERVAL", "6h")
if err != nil {
t.Fatal(err)
}
if got != 6*time.Hour {
t.Fatalf("durationFromEnvironment() = %s, want 6h", got)
}
t.Setenv("TEST_INTERVAL", "45m")
got, err = durationFromEnvironment("TEST_INTERVAL", "6h")
if err != nil {
t.Fatal(err)
}
if got != 45*time.Minute {
t.Fatalf("durationFromEnvironment() = %s, want 45m", got)
}
}
func TestDurationFromEnvironmentRejectsNonPositive(t *testing.T) {
t.Setenv("TEST_INTERVAL", "0s")
if _, err := durationFromEnvironment("TEST_INTERVAL", "6h"); err == nil {
t.Fatal("durationFromEnvironment() accepted a non-positive duration")
}
}
func TestDurationFromEnvironmentRejectsInvalidDuration(t *testing.T) {
t.Setenv("TEST_INTERVAL", "eventually")
if _, err := durationFromEnvironment("TEST_INTERVAL", "6h"); err == nil {
t.Fatal("durationFromEnvironment() accepted an invalid duration")
}
}
func TestProvidersFromEnvironment(t *testing.T) {
t.Setenv("ENABLED_PROVIDERS", "")
providers, err := providersFromEnvironment()
if err != nil {
t.Fatal(err)
}
got := mapsKeys(providers)
slices.Sort(got)
if want := []string{"google-vertex", "novita"}; !reflect.DeepEqual(got, want) {
t.Fatalf("default providers = %v, want %v", got, want)
}
t.Setenv("ENABLED_PROVIDERS", "openai, google-vertex")
providers, err = providersFromEnvironment()
if err != nil {
t.Fatal(err)
}
got = mapsKeys(providers)
slices.Sort(got)
if want := []string{"google-vertex", "openai"}; !reflect.DeepEqual(got, want) {
t.Fatalf("configured providers = %v, want %v", got, want)
}
}
func TestProvidersFromEnvironmentRejectsInvalidSelections(t *testing.T) {
for _, value := range []string{"unknown", "novita,novita", ","} {
t.Run(value, func(t *testing.T) {
t.Setenv("ENABLED_PROVIDERS", value)
if _, err := providersFromEnvironment(); err == nil {
t.Fatalf("providersFromEnvironment() accepted %q", value)
}
})
}
}
func mapsKeys[T any](input map[string]T) []string {
keys := make([]string, 0, len(input))
for key := range input {
keys = append(keys, key)
}
return keys
}
func selectProviders(names ...string) providerSelection {
definitions := buildProviderDefinitions()
providers := providerSelection{}
for _, name := range names {
providers[name] = definitions[name]
}
return providers
}
func number(value string) *json.Number {
number := json.Number(value)
return &number
}
func TestConvertModelsDev(t *testing.T) {
provider := modelsDevProvider{Models: map[string]modelsDevModel{
"model-a": {Cost: modelsDevCost{
Input: number("1.5"),
Output: number("7.5"),
CacheRead: number("0.15"),
Tiers: []modelsDevTier{{
Input: number("3"),
Output: number("15"),
Tier: modelsDevTierSelector{Type: "context", Size: 200000},
}},
}},
"unpriced": {},
}}
want := map[string]catalogModel{
"model-a": {
Rates: map[string]string{"input": "1.5", "output": "7.5", "cacheRead": "0.15"},
Tiers: []catalogTier{{ContextOver: 200000, Rates: map[string]string{"input": "3", "output": "15"}}},
},
}
if got := convertModelsDev(provider); !reflect.DeepEqual(got, want) {
t.Fatalf("convertModelsDev() = %#v, want %#v", got, want)
}
}
func TestConvertNovita(t *testing.T) {
models, err := convertNovita(novitaCatalog{Data: []novitaModel{{
ID: "model-a",
InputPrice: json.Number("3900"),
OutputPrice: json.Number("12500"),
}}})
if err != nil {
t.Fatal(err)
}
want := map[string]catalogModel{
"model-a": {Rates: map[string]string{"input": "0.39", "output": "1.25"}},
}
if !reflect.DeepEqual(models, want) {
t.Fatalf("convertNovita() = %#v, want %#v", models, want)
}
}
func TestConvertNovitaSkipsIncompleteModels(t *testing.T) {
models, err := convertNovita(novitaCatalog{Data: []novitaModel{
{ID: "", InputPrice: json.Number("10000"), OutputPrice: json.Number("20000")},
{ID: "missing-input", OutputPrice: json.Number("20000")},
{ID: "missing-output", InputPrice: json.Number("10000")},
}})
if err != nil {
t.Fatal(err)
}
if len(models) != 0 {
t.Fatalf("convertNovita() = %#v, want no models", models)
}
}
func TestBuildCatalogUsesAgentgatewayProviderIDs(t *testing.T) {
input := modelsDevCatalog{
"openai": {Models: map[string]modelsDevModel{"openai-model": {Cost: modelsDevCost{Input: number("1")}}}},
"google-vertex": {Models: map[string]modelsDevModel{
"vertex-model": {Cost: modelsDevCost{Input: number("2")}},
}},
}
output, err := buildCatalog(input, novitaCatalog{Data: []novitaModel{{
ID: "novita-model", InputPrice: json.Number("10000"), OutputPrice: json.Number("20000"),
}}}, selectProviders("openai", "google-vertex", "novita"))
if err != nil {
t.Fatal(err)
}
want := []string{"custom", "gcp.vertex_ai", "openai"}
got := make([]string, 0, len(output.Providers))
for provider := range output.Providers {
got = append(got, provider)
}
slices.Sort(got)
if !reflect.DeepEqual(got, want) {
t.Fatalf("provider IDs = %v, want %v", got, want)
}
}
func TestBuildCatalogRejectsMissingOrInvalidSources(t *testing.T) {
valid := modelsDevCatalog{
"google-vertex": {Models: map[string]modelsDevModel{
"vertex-model": {Cost: modelsDevCost{Input: number("2")}},
}},
}
tests := []struct {
name string
input modelsDevCatalog
novita novitaCatalog
}{
{
name: "missing models.dev provider",
input: modelsDevCatalog{},
novita: novitaCatalog{Data: []novitaModel{{
ID: "novita-model", InputPrice: json.Number("10000"), OutputPrice: json.Number("20000"),
}}},
},
{
name: "missing Novita models",
input: valid,
},
{
name: "invalid Novita price",
input: valid,
novita: novitaCatalog{Data: []novitaModel{{
ID: "novita-model", InputPrice: json.Number("-1"), OutputPrice: json.Number("20000"),
}}},
},
}
for _, test := range tests {
t.Run(test.name, func(t *testing.T) {
if _, err := buildCatalog(
test.input,
test.novita,
selectProviders("google-vertex", "novita"),
); err == nil {
t.Fatal("buildCatalog() unexpectedly succeeded")
}
})
}
}
func TestTenThousandths(t *testing.T) {
tests := []struct {
input string
want string
}{
{input: "0", want: "0"},
{input: "1", want: "0.0001"},
{input: "3900", want: "0.39"},
{input: "10000", want: "1"},
{input: "12500", want: "1.25"},
{input: "00012500", want: "1.25"},
}
for _, test := range tests {
t.Run(test.input, func(t *testing.T) {
got, err := tenThousandths(json.Number(test.input))
if err != nil {
t.Fatal(err)
}
if got != test.want {
t.Fatalf("tenThousandths(%q) = %q, want %q", test.input, got, test.want)
}
})
}
}
func TestTenThousandthsRejectsInvalidNumbers(t *testing.T) {
for _, input := range []string{"0.39", "-1", "1e4"} {
t.Run(input, func(t *testing.T) {
if _, err := tenThousandths(json.Number(input)); err == nil {
t.Fatalf("tenThousandths(%q) unexpectedly succeeded", input)
}
})
}
}
func TestRenderMetrics(t *testing.T) {
input := catalog{Providers: map[string]catalogProvider{
"custom": {Models: map[string]catalogModel{
"vendor/model": {Rates: map[string]string{"output": "1.25", "cacheRead": "0.1"}},
}},
}}
want := "# HELP agentgateway_model_token_price_usd_per_million Model token price in USD per million tokens.\n" +
"# TYPE agentgateway_model_token_price_usd_per_million gauge\n" +
"agentgateway_model_token_price_usd_per_million{gen_ai_system=\"custom\",gen_ai_request_model=\"vendor/model\",gen_ai_token_type=\"cache_read\",source_provider=\"novita\"} 0.1\n" +
"agentgateway_model_token_price_usd_per_million{gen_ai_system=\"custom\",gen_ai_request_model=\"vendor/model\",gen_ai_token_type=\"output\",source_provider=\"novita\"} 1.25\n"
if got := string(renderMetrics(input)); got != want {
t.Fatalf("renderMetrics() = %q, want %q", got, want)
}
}
func TestRenderModelList(t *testing.T) {
input := catalog{Providers: map[string]catalogProvider{
"gcp.vertex_ai": {Models: map[string]catalogModel{"gemini-z": {}, "gemini-a": {}}},
"custom": {Models: map[string]catalogModel{"vendor/model": {}}},
}}
want := modelList{Object: "list", Data: []modelListItem{
{ID: "novita/vendor/model", Object: "model", Created: 0, OwnedBy: "novita"},
{ID: "vertex_ai/gemini-a", Object: "model", Created: 0, OwnedBy: "google-vertex"},
{ID: "vertex_ai/gemini-z", Object: "model", Created: 0, OwnedBy: "google-vertex"},
}}
if got := renderModelList(input); !reflect.DeepEqual(got, want) {
t.Fatalf("renderModelList() = %#v, want %#v", got, want)
}
}
func TestReadCatalog(t *testing.T) {
path := filepath.Join(t.TempDir(), "catalog.json")
want := catalog{Providers: map[string]catalogProvider{
"gcp.vertex_ai": {Models: map[string]catalogModel{"gemini": {Rates: map[string]string{"input": "1"}}}},
}}
contents, err := json.Marshal(want)
if err != nil {
t.Fatal(err)
}
if err := os.WriteFile(path, contents, 0o600); err != nil {
t.Fatal(err)
}
got, err := readCatalog(path)
if err != nil {
t.Fatal(err)
}
if !reflect.DeepEqual(got, want) {
t.Fatalf("readCatalog() = %#v, want %#v", got, want)
}
if err := os.WriteFile(path, []byte("{"), 0o600); err != nil {
t.Fatal(err)
}
if _, err := readCatalog(path); err == nil {
t.Fatal("readCatalog() accepted malformed JSON")
}
if _, err := readCatalog(filepath.Join(t.TempDir(), "missing")); err == nil {
t.Fatal("readCatalog() accepted a missing file")
}
}
func TestCatalogHandler(t *testing.T) {
path := filepath.Join(t.TempDir(), "catalog.json")
contents := `{"providers":{
"gcp.vertex_ai":{"models":{"gemini":{"rates":{"input":"2"}}}},
"custom":{"models":{"novita/model":{"rates":{"input":"3"}}}}
}}`
if err := os.WriteFile(path, []byte(contents), 0o600); err != nil {
t.Fatal(err)
}
handler := catalogHandler(path, selectProviders("google-vertex", "novita"))
tests := []struct {
name string
method string
path string
status int
contentType string
body string
}{
{name: "health", method: http.MethodGet, path: "/healthz", status: http.StatusNoContent},
{name: "ready", method: http.MethodGet, path: "/readyz", status: http.StatusNoContent},
{
name: "metrics",
method: http.MethodGet,
path: "/metrics",
status: http.StatusOK,
contentType: "text/plain; version=0.0.4",
body: `gen_ai_request_model="gemini"`,
},
{
name: "models",
method: http.MethodGet,
path: "/v1/models",
status: http.StatusOK,
contentType: "application/json",
body: `"id":"vertex_ai/gemini"`,
},
{name: "models method", method: http.MethodPost, path: "/v1/models", status: http.StatusMethodNotAllowed},
}
for _, test := range tests {
t.Run(test.name, func(t *testing.T) {
request := httptest.NewRequest(test.method, test.path, nil)
response := httptest.NewRecorder()
handler.ServeHTTP(response, request)
if response.Code != test.status {
t.Fatalf("status = %d, want %d", response.Code, test.status)
}
if test.contentType != "" && response.Header().Get("Content-Type") != test.contentType {
t.Fatalf("Content-Type = %q, want %q", response.Header().Get("Content-Type"), test.contentType)
}
if test.body != "" && !strings.Contains(response.Body.String(), test.body) {
t.Fatalf("body = %q, want substring %q", response.Body.String(), test.body)
}
})
}
}
func TestCatalogHandlerRejectsInvalidCatalog(t *testing.T) {
missing := filepath.Join(t.TempDir(), "missing")
providers := selectProviders("google-vertex", "novita")
handler := catalogHandler(missing, providers)
for _, path := range []string{"/readyz", "/metrics", "/v1/models"} {
t.Run(path, func(t *testing.T) {
request := httptest.NewRequest(http.MethodGet, path, nil)
response := httptest.NewRecorder()
handler.ServeHTTP(response, request)
if response.Code != http.StatusServiceUnavailable {
t.Fatalf("status = %d, want %d", response.Code, http.StatusServiceUnavailable)
}
})
}
empty := filepath.Join(t.TempDir(), "catalog.json")
if err := os.WriteFile(empty, []byte(`{"providers":{}}`), 0o600); err != nil {
t.Fatal(err)
}
request := httptest.NewRequest(http.MethodGet, "/readyz", nil)
response := httptest.NewRecorder()
catalogHandler(empty, providers).ServeHTTP(response, request)
if response.Code != http.StatusServiceUnavailable {
t.Fatalf("empty catalog status = %d, want %d", response.Code, http.StatusServiceUnavailable)
}
}
func TestFetchJSON(t *testing.T) {
client := &http.Client{Transport: roundTripFunc(func(request *http.Request) (*http.Response, error) {
if request.Header.Get("Accept") != "application/json" {
t.Errorf("Accept = %q", request.Header.Get("Accept"))
}
if request.Header.Get("Authorization") != "Bearer test" {
t.Errorf("Authorization = %q", request.Header.Get("Authorization"))
}
return response(http.StatusOK, `{"value":12345678901234567890}`), nil
})}
var target struct {
Value json.Number `json:"value"`
}
if err := fetchJSON(context.Background(), client, "https://models.invalid", "Bearer test", &target); err != nil {
t.Fatal(err)
}
if target.Value.String() != "12345678901234567890" {
t.Fatalf("value = %q", target.Value)
}
}
func TestFetchJSONRejectsHTTPAndJSONErrors(t *testing.T) {
tests := []struct {
name string
status int
body string
}{
{name: "status", status: http.StatusBadGateway, body: `{"error":"upstream"}`},
{name: "JSON", status: http.StatusOK, body: "{"},
}
for _, test := range tests {
t.Run(test.name, func(t *testing.T) {
client := &http.Client{Transport: roundTripFunc(func(*http.Request) (*http.Response, error) {
return response(test.status, test.body), nil
})}
if err := fetchJSON(context.Background(), client, "https://models.invalid", "", &map[string]any{}); err == nil {
t.Fatal("fetchJSON() unexpectedly succeeded")
}
})
}
}
func TestWriteConfigMapPatchesExistingConfigMap(t *testing.T) {
client := &http.Client{Transport: roundTripFunc(func(request *http.Request) (*http.Response, error) {
if request.Method != http.MethodPatch {
t.Errorf("method = %q", request.Method)
}
if request.URL.Path != "/configmaps/agentgateway-model-costs" {
t.Errorf("path = %q", request.URL.Path)
}
if request.Header.Get("Authorization") != "Bearer token" {
t.Errorf("Authorization = %q", request.Header.Get("Authorization"))
}
if request.Header.Get("Content-Type") != "application/merge-patch+json" {
t.Errorf("Content-Type = %q", request.Header.Get("Content-Type"))
}
var patch configMapPatch
if err := json.NewDecoder(request.Body).Decode(&patch); err != nil {
t.Error(err)
}
if patch.Data["catalog.json"] != `{"providers":{}}` {
t.Errorf("catalog.json = %q", patch.Data["catalog.json"])
}
return response(http.StatusOK, `{}`), nil
})}
err := writeConfigMap(
context.Background(),
client,
"https://kubernetes.invalid/configmaps",
[]byte("token"),
[]byte(`{"providers":{}}`),
)
if err != nil {
t.Fatal(err)
}
}
func TestWriteConfigMapCreatesMissingConfigMap(t *testing.T) {
requests := 0
client := &http.Client{Transport: roundTripFunc(func(request *http.Request) (*http.Response, error) {
requests++
if requests == 1 {
if request.Method != http.MethodPatch {
t.Errorf("first method = %q", request.Method)
}
return response(http.StatusNotFound, ""), nil
}
if request.Method != http.MethodPost {
t.Errorf("second method = %q", request.Method)
}
if request.URL.Path != "/configmaps" {
t.Errorf("path = %q", request.URL.Path)
}
var resource configMapResource
if err := json.NewDecoder(request.Body).Decode(&resource); err != nil {
t.Error(err)
}
if resource.APIVersion != "v1" || resource.Kind != "ConfigMap" {
t.Errorf("resource = %#v", resource)
}
if resource.Metadata.Name != configMapName {
t.Errorf("name = %q", resource.Metadata.Name)
}
if resource.Metadata.Annotations["kustomize.toolkit.fluxcd.io/prune"] != "disabled" {
t.Errorf("annotations = %#v", resource.Metadata.Annotations)
}
if resource.Metadata.Labels["app.kubernetes.io/instance"] != "agentgateway" ||
resource.Metadata.Labels["app.kubernetes.io/name"] != "model-catalog" ||
resource.Metadata.Labels["app.kubernetes.io/part-of"] != "agentgateway" {
t.Errorf("labels = %#v", resource.Metadata.Labels)
}
if resource.Data["catalog.json"] != `{"providers":{}}` {
t.Errorf("catalog.json = %q", resource.Data["catalog.json"])
}
return response(http.StatusCreated, `{}`), nil
})}
if err := writeConfigMap(
context.Background(),
client,
"https://kubernetes.invalid/configmaps",
[]byte("token"),
[]byte(`{"providers":{}}`),
); err != nil {
t.Fatal(err)
}
if requests != 2 {
t.Fatalf("requests = %d, want 2", requests)
}
}
func TestWriteConfigMapRetriesPatchAfterCreateConflict(t *testing.T) {
requests := 0
client := &http.Client{Transport: roundTripFunc(func(request *http.Request) (*http.Response, error) {
requests++
switch requests {
case 1:
return response(http.StatusNotFound, ""), nil
case 2:
return response(http.StatusConflict, ""), nil
default:
if request.Method != http.MethodPatch {
t.Errorf("third method = %q", request.Method)
}
return response(http.StatusOK, `{}`), nil
}
})}
if err := writeConfigMap(
context.Background(),
client,
"https://kubernetes.invalid/configmaps",
[]byte("token"),
[]byte(`{"providers":{}}`),
); err != nil {
t.Fatal(err)
}
if requests != 3 {
t.Fatalf("requests = %d, want 3", requests)
}
}
func TestWriteConfigMapRejectsHTTPAndTransportErrors(t *testing.T) {
client := &http.Client{Transport: roundTripFunc(func(*http.Request) (*http.Response, error) {
return response(http.StatusForbidden, ""), nil
})}
if err := writeConfigMap(context.Background(), client, "https://kubernetes.invalid", []byte("token"), []byte("{}")); err == nil {
t.Fatal("writeConfigMap() accepted an HTTP error")
}
client = &http.Client{Transport: roundTripFunc(func(*http.Request) (*http.Response, error) {
return nil, errors.New("transport failed")
})}
if err := writeConfigMap(context.Background(), client, "https://kubernetes.invalid", []byte("token"), []byte("{}")); err == nil {
t.Fatal("writeConfigMap() accepted a transport error")
}
}
func TestKubernetesRequestDecodesResponse(t *testing.T) {
client := &http.Client{Transport: roundTripFunc(func(request *http.Request) (*http.Response, error) {
if request.Header.Get("Content-Type") != "" {
t.Errorf("Content-Type = %q, want empty", request.Header.Get("Content-Type"))
}
return response(http.StatusOK, `{"value":"decoded"}`), nil
})}
target := map[string]string{}
status, err := kubernetesRequest(
context.Background(),
client,
http.MethodGet,
"https://kubernetes.invalid",
[]byte("token"),
nil,
&target,
)
if err != nil {
t.Fatal(err)
}
if status != http.StatusOK || target["value"] != "decoded" {
t.Fatalf("status = %d, target = %#v", status, target)
}
}
type roundTripFunc func(*http.Request) (*http.Response, error)
func (function roundTripFunc) RoundTrip(request *http.Request) (*http.Response, error) {
return function(request)
}
func response(status int, body string) *http.Response {
return &http.Response{
StatusCode: status,
Status: http.StatusText(status),
Header: make(http.Header),
Body: io.NopCloser(strings.NewReader(body)),
}
}