1package provider
2
3import (
4 "context"
5 "errors"
6 "fmt"
7 "io"
8 "net/http"
9 "net/http/httptest"
10 "strings"
11 "testing"
12)
13
14func TestCustomModelsURL(t *testing.T) {
15 cases := map[string]string{
16 "http://127.0.0.1:48000": "http://127.0.0.1:48000/v1/models",
17 "http://127.0.0.1:48000/": "http://127.0.0.1:48000/v1/models",
18 "http://127.0.0.1:48000/v1": "http://127.0.0.1:48000/v1/models",
19 "http://127.0.0.1:48000/v1/": "http://127.0.0.1:48000/v1/models",
20 "https://llm.example.com/paas/v4": "https://llm.example.com/paas/v4/models",
21 }
22 for in, want := range cases {
23 got, err := customModelsURL(in)
24 if err != nil || got != want {
25 t.Errorf("customModelsURL(%q) = %q, %v; want %q", in, got, err, want)
26 }
27 }
28 for _, bad := range []string{"", "localhost:1234", "ftp://x", "/v1"} {
29 if _, err := customModelsURL(bad); err == nil {
30 t.Errorf("customModelsURL(%q) accepted", bad)
31 }
32 }
33}
34
35func TestDiscoverCustomProvider(t *testing.T) {
36 server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
37 if r.URL.Path != "/v1/models" {
38 t.Errorf("unexpected path %s", r.URL.Path)
39 }
40 if r.Header.Get("Authorization") != "Bearer dummy" {
41 t.Errorf("missing bearer header: %q", r.Header.Get("Authorization"))
42 }
43 fmt.Fprint(w, `{"object":"list","data":[{"id":"qwen3-coder"},{"id":"nomic-embed-text"},{"id":"qwen3-coder"},{"id":""},{"id":"gpt-oss-20b"}]}`)
44 }))
45 defer server.Close()
46
47 models, err := DiscoverCustomProvider(context.Background(), "m4", server.URL+"/v1", "dummy")
48 if err != nil {
49 t.Fatal(err)
50 }
51 if len(models) != 2 {
52 t.Fatalf("models = %+v", models)
53 }
54 first := models[0]
55 if first.Provider != "m4" || first.ID != "qwen3-coder" || first.Source != "live" || first.BaseURL != server.URL+"/v1" {
56 t.Fatalf("model = %+v", first)
57 }
58 if first.ContextWindow != customDiscoveryContextWindow || first.MaxOutput != customDiscoveryMaxOutput {
59 t.Fatalf("defaults = %d/%d", first.ContextWindow, first.MaxOutput)
60 }
61 if models[1].ID != "gpt-oss-20b" {
62 t.Fatalf("second model = %+v", models[1])
63 }
64}
65
66func TestDiscoverCustomProviderOmitsHeaderWithoutKey(t *testing.T) {
67 server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
68 if _, ok := r.Header["Authorization"]; ok {
69 t.Error("unexpected Authorization header")
70 }
71 fmt.Fprint(w, `{"data":[{"id":"local"}]}`)
72 }))
73 defer server.Close()
74 models, err := DiscoverCustomProvider(context.Background(), "m5", server.URL, "")
75 if err != nil || len(models) != 1 {
76 t.Fatalf("models=%+v err=%v", models, err)
77 }
78}
79
80func TestDiscoverCustomProviderErrors(t *testing.T) {
81 status := http.StatusUnauthorized
82 body := `{}`
83 server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
84 w.WriteHeader(status)
85 fmt.Fprint(w, body)
86 }))
87 defer server.Close()
88
89 if _, err := DiscoverCustomProvider(context.Background(), "m4", server.URL, "k"); err == nil || err.Error() != "m4 discovery HTTP 401" {
90 t.Fatalf("status error = %v", err)
91 }
92 status = http.StatusOK
93 body = `{"object":"list"}`
94 if _, err := DiscoverCustomProvider(context.Background(), "m4", server.URL, "k"); err == nil {
95 t.Fatal("missing data accepted")
96 }
97 body = `not json`
98 if _, err := DiscoverCustomProvider(context.Background(), "m4", server.URL, "k"); err == nil {
99 t.Fatal("invalid JSON accepted")
100 }
101 if _, err := DiscoverCustomProvider(context.Background(), "", server.URL, "k"); err == nil {
102 t.Fatal("empty provider id accepted")
103 }
104
105 ctx, cancel := context.WithCancel(context.Background())
106 cancel()
107 if _, err := DiscoverCustomProvider(ctx, "m4", server.URL, "k"); !errors.Is(err, context.Canceled) {
108 t.Fatalf("cancellation error = %v", err)
109 }
110}
111
112func TestDiscoverCustomProviderBoundsResponseAndUsesClient(t *testing.T) {
113 big := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
114 fmt.Fprint(w, `{"data":[{"id":"`)
115 chunk := make([]byte, 1<<20)
116 for i := range chunk {
117 chunk[i] = 'a'
118 }
119 for i := 0; i < 9; i++ {
120 w.Write(chunk)
121 }
122 fmt.Fprint(w, `"}]}`)
123 }))
124 defer big.Close()
125 if _, err := DiscoverCustomProvider(context.Background(), "m4", big.URL, ""); err == nil {
126 t.Fatal("oversized response accepted")
127 }
128
129 used := false
130 client := &http.Client{Transport: roundTripFunc(func(r *http.Request) (*http.Response, error) {
131 used = true
132 return &http.Response{StatusCode: 200, Body: io.NopCloser(strings.NewReader(`{"data":[{"id":"x"}]}`)), Header: http.Header{}}, nil
133 })}
134 models, err := DiscoverCustomProviderWithClient(context.Background(), client, "m4", "http://example.invalid/v1", "")
135 if err != nil || !used || len(models) != 1 {
136 t.Fatalf("custom client not used: used=%v models=%v err=%v", used, models, err)
137 }
138}