chore: initial commit
This commit is contained in:
@@ -0,0 +1,273 @@
|
||||
package server
|
||||
|
||||
import (
|
||||
"context"
|
||||
"encoding/json"
|
||||
"fmt"
|
||||
"net/http"
|
||||
"net/http/httptest"
|
||||
"reflect"
|
||||
"strings"
|
||||
"testing"
|
||||
"time"
|
||||
)
|
||||
|
||||
const discoveryPath = "/admin/api/channels/discover-models"
|
||||
|
||||
func discover(t *testing.T, s *Server, cookie *http.Cookie, c channel) *httptest.ResponseRecorder {
|
||||
t.Helper()
|
||||
b, e := json.Marshal(c)
|
||||
if e != nil {
|
||||
t.Fatal(e)
|
||||
}
|
||||
return call(s, "POST", discoveryPath, string(b), cookie)
|
||||
}
|
||||
|
||||
func TestDiscoveryProtocolsAndReadOnly(t *testing.T) {
|
||||
for _, proto := range []string{"openai-chat", "openai-responses", "anthropic"} {
|
||||
t.Run(proto, func(t *testing.T) {
|
||||
s := testServer(t)
|
||||
cookie := loginCookie(t, s)
|
||||
calls := 0
|
||||
up := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||
calls++
|
||||
if r.Method != "GET" || r.URL.Path != "/v1/models" {
|
||||
t.Error(r.Method, r.URL)
|
||||
}
|
||||
if proto == "anthropic" {
|
||||
if r.Header.Get("x-api-key") != "secret" || r.Header.Get("anthropic-version") != "2023-06-01" || r.Header.Get("Authorization") != "" {
|
||||
t.Error("wrong anthropic headers")
|
||||
}
|
||||
if r.URL.Query().Get("after_id") == "" {
|
||||
fmt.Fprint(w, `{"data":[{"id":"z"}],"has_more":true,"last_id":"cursor/&?"}`)
|
||||
return
|
||||
}
|
||||
if r.URL.Query().Get("after_id") != "cursor/&?" {
|
||||
t.Error(r.URL)
|
||||
}
|
||||
} else if r.Header.Get("Authorization") != "Bearer secret" || r.Header.Get("x-api-key") != "" {
|
||||
t.Error("wrong openai headers")
|
||||
}
|
||||
fmt.Fprint(w, `{"data":[{"id":"z"},{"id":"a"},{"id":"a"}]}`)
|
||||
}))
|
||||
defer up.Close()
|
||||
draft := channel{BaseURL: up.URL + "/v1/", Protocol: proto, APIKey: "secret"}
|
||||
if proto == "anthropic" {
|
||||
draft.BaseURL = up.URL + "/"
|
||||
}
|
||||
w := discover(t, s, cookie, draft)
|
||||
if w.Code != 200 || strings.TrimSpace(w.Body.String()) != `{"models":["a","z"]}` {
|
||||
t.Fatal(w.Code, w.Body.String())
|
||||
}
|
||||
wantCalls := 1
|
||||
if proto == "anthropic" {
|
||||
wantCalls = 2
|
||||
}
|
||||
if calls != wantCalls {
|
||||
t.Fatal(calls)
|
||||
}
|
||||
cs, e := s.loadChannels()
|
||||
if e != nil || len(cs) != 0 {
|
||||
t.Fatal(cs, e)
|
||||
}
|
||||
draft.Name = "saved"
|
||||
draft.Models = map[string]string{"alias": "a"}
|
||||
b, _ := json.Marshal(draft)
|
||||
w = call(s, "POST", "/admin/api/channels", string(b), cookie)
|
||||
if w.Code != 200 {
|
||||
t.Fatal(w.Body.String())
|
||||
}
|
||||
var saved channel
|
||||
if e := json.Unmarshal(w.Body.Bytes(), &saved); e != nil {
|
||||
t.Fatal(e)
|
||||
}
|
||||
before, _ := s.loadChannels()
|
||||
draft.ID, draft.APIKey, draft.Name, draft.Models = saved.ID, "", "unsaved", nil
|
||||
if w = discover(t, s, cookie, draft); w.Code != 200 {
|
||||
t.Fatal(w.Code, w.Body.String())
|
||||
}
|
||||
after, _ := s.loadChannels()
|
||||
if !reflect.DeepEqual(before, after) {
|
||||
t.Fatal("discovery changed channel")
|
||||
}
|
||||
replacement := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||
if proto == "anthropic" {
|
||||
if r.Header.Get("x-api-key") != "replacement" {
|
||||
t.Error("draft key not used")
|
||||
}
|
||||
} else if r.Header.Get("Authorization") != "Bearer replacement" {
|
||||
t.Error("draft key not used")
|
||||
}
|
||||
fmt.Fprint(w, `{"data":[]}`)
|
||||
}))
|
||||
defer replacement.Close()
|
||||
draft.BaseURL, draft.APIKey = replacement.URL, "replacement"
|
||||
w = discover(t, s, cookie, draft)
|
||||
if w.Code != 200 || strings.TrimSpace(w.Body.String()) != `{"models":[]}` {
|
||||
t.Fatal(w.Code, w.Body.String())
|
||||
}
|
||||
after, _ = s.loadChannels()
|
||||
if !reflect.DeepEqual(before, after) {
|
||||
t.Fatal("draft replacement was saved")
|
||||
}
|
||||
var logs int
|
||||
if e := s.db.QueryRow("SELECT count(*) FROM logs").Scan(&logs); e != nil || logs != 0 {
|
||||
t.Fatal(logs, e)
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
func TestDiscoveryAuthAndValidation(t *testing.T) {
|
||||
s := testServer(t)
|
||||
cookie := loginCookie(t, s)
|
||||
if w := call(s, "POST", discoveryPath, "{}", nil); w.Code != 401 {
|
||||
t.Fatal(w.Code)
|
||||
}
|
||||
for _, header := range []string{"Origin", "X-Requested-With"} {
|
||||
r := httptest.NewRequest("POST", discoveryPath, strings.NewReader("{}"))
|
||||
r.AddCookie(cookie)
|
||||
r.Header.Set("Origin", "http://example.com")
|
||||
r.Header.Set("X-Requested-With", "ai-router")
|
||||
r.Header.Del(header)
|
||||
w := httptest.NewRecorder()
|
||||
s.ServeHTTP(w, r)
|
||||
if w.Code != 403 {
|
||||
t.Fatal(header, w.Code)
|
||||
}
|
||||
}
|
||||
for _, method := range []string{"GET", "DELETE", "PUT"} {
|
||||
if w := call(s, method, discoveryPath, "{}", cookie); w.Code != 405 {
|
||||
t.Fatal(method, w.Code)
|
||||
}
|
||||
}
|
||||
for _, base := range []string{"", "ftp://example.com", "http://", "http://user:[email protected]", "http://example.com?x=1", "http://example.com?", "http://example.com#"} {
|
||||
if w := discover(t, s, cookie, channel{BaseURL: base, Protocol: "openai-chat"}); w.Code != 400 {
|
||||
t.Fatal(base, w.Code)
|
||||
}
|
||||
}
|
||||
if w := discover(t, s, cookie, channel{BaseURL: "http://example.com", Protocol: "unknown"}); w.Code != 400 {
|
||||
t.Fatal(w.Code)
|
||||
}
|
||||
if w := discover(t, s, cookie, channel{ID: "missing", BaseURL: "http://example.com", Protocol: "anthropic"}); w.Code != 404 {
|
||||
t.Fatal(w.Code)
|
||||
}
|
||||
for _, body := range []string{"{", "{} {}", strings.Repeat(" ", maxBody+1)} {
|
||||
if w := call(s, "POST", discoveryPath, body, cookie); w.Code != 400 {
|
||||
t.Fatal(w.Code)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func TestDiscoveryUpstreamErrors(t *testing.T) {
|
||||
for _, tc := range []struct {
|
||||
name, body string
|
||||
status, calls int
|
||||
}{
|
||||
{"status", "secret upstream details", 401, 1},
|
||||
{"json", "secret", 200, 1},
|
||||
{"missing data", "{}", 200, 1},
|
||||
{"null", `{"data":null}`, 200, 1},
|
||||
{"empty id", `{"data":[{"id":""}]}`, 200, 1},
|
||||
{"trailing", `{"data":[]}{}`, 200, 1},
|
||||
{"oversized", strings.Repeat("x", discoveryMaxBody+1), 200, 1},
|
||||
{"missing cursor", `{"data":[],"has_more":true}`, 200, 1},
|
||||
{"repeated cursor", `{"data":[],"has_more":true,"last_id":"same"}`, 200, 2},
|
||||
{"page limit", "", 200, discoveryMaxPages},
|
||||
} {
|
||||
t.Run(tc.name, func(t *testing.T) {
|
||||
s := testServer(t)
|
||||
cookie := loginCookie(t, s)
|
||||
calls := 0
|
||||
up := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||
calls++
|
||||
w.WriteHeader(tc.status)
|
||||
if tc.name == "page limit" {
|
||||
fmt.Fprintf(w, `{"data":[],"has_more":true,"last_id":"%d"}`, calls)
|
||||
} else {
|
||||
fmt.Fprint(w, tc.body)
|
||||
}
|
||||
}))
|
||||
defer up.Close()
|
||||
w := discover(t, s, cookie, channel{BaseURL: up.URL, Protocol: "anthropic", APIKey: "secret"})
|
||||
if w.Code != 502 || strings.TrimSpace(w.Body.String()) != `{"error":"model discovery failed"}` || calls != tc.calls {
|
||||
t.Fatal(w.Code, w.Body.String(), calls)
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
func TestDiscoveryRedirectAndDeadline(t *testing.T) {
|
||||
s := testServer(t)
|
||||
cookie := loginCookie(t, s)
|
||||
target := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { t.Error("redirect followed") }))
|
||||
defer target.Close()
|
||||
up := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { http.Redirect(w, r, target.URL, 307) }))
|
||||
defer up.Close()
|
||||
for _, proto := range []string{"openai-chat", "anthropic"} {
|
||||
if w := discover(t, s, cookie, channel{BaseURL: up.URL, Protocol: proto, APIKey: "secret"}); w.Code != 502 {
|
||||
t.Fatal(w.Code)
|
||||
}
|
||||
}
|
||||
slow := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { <-r.Context().Done() }))
|
||||
defer slow.Close()
|
||||
ctx, cancel := context.WithTimeout(context.Background(), 20*time.Millisecond)
|
||||
defer cancel()
|
||||
b, _ := json.Marshal(channel{BaseURL: slow.URL, Protocol: "openai-chat"})
|
||||
r := httptest.NewRequest("POST", discoveryPath, strings.NewReader(string(b))).WithContext(ctx)
|
||||
r.AddCookie(cookie)
|
||||
r.Header.Set("Origin", "http://example.com")
|
||||
r.Header.Set("X-Requested-With", "ai-router")
|
||||
w := httptest.NewRecorder()
|
||||
s.ServeHTTP(w, r)
|
||||
if w.Code != 504 {
|
||||
t.Fatal(w.Code, w.Body.String())
|
||||
}
|
||||
}
|
||||
|
||||
func TestModelsListSelectedAliasesAndPermissions(t *testing.T) {
|
||||
s := testServer(t)
|
||||
cookie := loginCookie(t, s)
|
||||
var allowed string
|
||||
for _, name := range []string{"selected", "forbidden", "disabled"} {
|
||||
b, _ := json.Marshal(channel{Name: name, BaseURL: "http://example.com", Protocol: "openai-chat", Enabled: name != "disabled", Models: map[string]string{"z": "vendor/real", "alias": "vendor/real"}})
|
||||
w := call(s, "POST", "/admin/api/channels", string(b), cookie)
|
||||
if w.Code != 200 {
|
||||
t.Fatal(w.Body.String())
|
||||
}
|
||||
var c channel
|
||||
json.Unmarshal(w.Body.Bytes(), &c)
|
||||
if name != "forbidden" {
|
||||
if allowed != "" {
|
||||
allowed += ","
|
||||
}
|
||||
allowed += fmt.Sprintf("%q", c.ID)
|
||||
}
|
||||
}
|
||||
w := call(s, "POST", "/admin/api/keys", `{"name":"client","channels":[`+allowed+`]}`, cookie)
|
||||
if w.Code != 200 {
|
||||
t.Fatal(w.Body.String())
|
||||
}
|
||||
var key map[string]string
|
||||
json.Unmarshal(w.Body.Bytes(), &key)
|
||||
r := httptest.NewRequest("GET", "/v1/models", nil)
|
||||
r.Header.Set("Authorization", "Bearer "+key["key"])
|
||||
w = httptest.NewRecorder()
|
||||
s.ServeHTTP(w, r)
|
||||
var result struct{ Data []struct{ ID string } }
|
||||
if e := json.Unmarshal(w.Body.Bytes(), &result); e != nil || w.Code != 200 {
|
||||
t.Fatal(w.Code, e)
|
||||
}
|
||||
if len(result.Data) != 2 || result.Data[0].ID != "selected/alias" || result.Data[1].ID != "selected/z" {
|
||||
t.Fatal(w.Body.String())
|
||||
}
|
||||
for _, model := range []string{"forbidden/alias", "disabled/alias"} {
|
||||
r = httptest.NewRequest("POST", "/v1/chat/completions", strings.NewReader(fmt.Sprintf(`{"model":%q,"messages":[]}`, model)))
|
||||
r.Header.Set("Authorization", "Bearer "+key["key"])
|
||||
w = httptest.NewRecorder()
|
||||
s.ServeHTTP(w, r)
|
||||
if w.Code != 403 {
|
||||
t.Fatal(model, w.Code)
|
||||
}
|
||||
}
|
||||
}
|
||||
Reference in New Issue
Block a user