1package agent
2
3import (
4 "fmt"
5 "io"
6 "net/http"
7 "strings"
8 "testing"
9 "time"
10
11 "github.com/patriceckhart/zot/packages/provider"
12)
13
14type yoloAutoTestTransport func(*http.Request) (*http.Response, error)
15
16func (f yoloAutoTestTransport) RoundTrip(r *http.Request) (*http.Response, error) {
17 return f(r)
18}
19
20func TestYoloAutoDefaultAvailableOffline(t *testing.T) {
21 id := defaultModelForProvider("yolo-auto")
22 if id != "qwen3.8-flash" {
23 t.Fatalf("default = %q, want free-plan model", id)
24 }
25 if _, err := provider.FindModel("yolo-auto", id); err != nil {
26 t.Fatal(err)
27 }
28}
29
30func TestYoloAutoRefreshIgnoresSharedCache(t *testing.T) {
31 t.Setenv("ZOT_HOME", t.TempDir())
32 t.Setenv("OPENROUTER_API_KEY", "")
33 t.Setenv("YOLO_AUTO_API_KEY", "test-pro")
34 cachedAt := time.Now().UTC()
35 if err := provider.SaveCache(ModelCachePath(), provider.ModelCache{
36 FetchedAt: cachedAt,
37 Models: []provider.Model{
38 {Provider: "openai", ID: "cached-public-model"},
39 {Provider: "yolo-auto", ID: "old-account-model", ContextWindow: 262144},
40 },
41 }); err != nil {
42 t.Fatal(err)
43 }
44 provider.SetLiveModels(nil)
45 t.Cleanup(func() { provider.SetLiveModels(nil) })
46 LoadCachedModels()
47
48 original := http.DefaultTransport
49 t.Cleanup(func() { http.DefaultTransport = original })
50 calls := 0
51 http.DefaultTransport = yoloAutoTestTransport(func(r *http.Request) (*http.Response, error) {
52 if r.URL.String() != "https://yolo-auto.com/v1/models" {
53 return nil, fmt.Errorf("unexpected discovery request: %s", r.URL)
54 }
55 calls++
56 body := `{"data":[{"id":"qwen3.8-flash","context_length":262144,"thinking":["high"]},{"id":"pro-only"}]}`
57 switch r.Header.Get("Authorization") {
58 case "Bearer test-pro":
59 case "Bearer test-free":
60 body = `{"data":[{"id":"qwen3.8-flash","context_length":131072,"thinking":["high"]}]}`
61 default:
62 return nil, fmt.Errorf("unexpected credential")
63 }
64 return &http.Response{StatusCode: http.StatusOK, Body: io.NopCloser(strings.NewReader(body)), Header: make(http.Header)}, nil
65 })
66
67 for _, tc := range []struct {
68 key string
69 window int
70 }{
71 {"test-pro", 262144},
72 {"test-free", 131072},
73 } {
74 t.Setenv("YOLO_AUTO_API_KEY", tc.key)
75 refreshModels()
76 model, err := provider.FindModel("yolo-auto", "qwen3.8-flash")
77 if err != nil {
78 t.Fatal(err)
79 }
80 if model.ContextWindow != tc.window {
81 t.Errorf("context = %d, want %d", model.ContextWindow, tc.window)
82 }
83 }
84 if calls != 2 {
85 t.Errorf("discovery calls = %d, want 2 despite fresh cache", calls)
86 }
87 if _, err := provider.FindModel("yolo-auto", "pro-only"); err == nil {
88 t.Error("previous account model survived refresh")
89 }
90 cached, err := provider.LoadCache(ModelCachePath())
91 if err != nil {
92 t.Fatal(err)
93 }
94 if !cached.FetchedAt.Equal(cachedAt) {
95 t.Error("account refresh extended public cache lifetime")
96 }
97 for _, model := range cached.Models {
98 if model.Provider == "yolo-auto" && (model.ID != "qwen3.8-flash" || model.ContextWindow != 131072) {
99 t.Errorf("previous account metadata survived in cache: %+v", model)
100 }
101 }
102 if _, err := provider.FindModel("yolo-auto", "old-account-model"); err == nil {
103 t.Error("previous cached account model survived refresh")
104 }
105 if _, err := provider.FindModel("openai", "cached-public-model"); err != nil {
106 t.Error("unrelated cached model lost")
107 }
108
109 // Failed discovery and logout must drop the previous account's metadata.
110 for _, key := range []string{"test-invalid", ""} {
111 t.Setenv("YOLO_AUTO_API_KEY", "test-pro")
112 refreshModels()
113 t.Setenv("YOLO_AUTO_API_KEY", key)
114 refreshModels()
115 model, err := provider.FindModel("yolo-auto", "qwen3.8-flash")
116 if err != nil {
117 t.Fatal(err)
118 }
119 if model.Source != "catalog" || len(model.ReasoningLevelMap) != 0 {
120 t.Errorf("account metadata survived failed discovery or logout: %+v", model)
121 }
122 if _, err := provider.FindModel("yolo-auto", "pro-only"); err == nil {
123 t.Error("previous account model survived failed discovery or logout")
124 }
125 }
126}