git.iko.soy/crush
on main help sign in
refactor: simplify context value retrieval using generics
07d065dwanghuaiyu@qiniu.com
internal/agent/tools/context_test.gonew file+219
compare with main
unifiedsplit
@@ -0,0 +1,219 @@
1+package tools
2+
3+import (
4+ "context"
5+ "testing"
6+)
7+
8+func TestGetContextValue(t *testing.T) {
9+ tests := []struct {
10+ name string
11+ setup func(ctx context.Context) context.Context
12+ key any
13+ defaultValue any
14+ want any
15+ }{
16+ {
17+ name: "returns string value",
18+ setup: func(ctx context.Context) context.Context {
19+ return context.WithValue(ctx, "testKey", "testValue")
20+ },
21+ key: "testKey",
22+ defaultValue: "",
23+ want: "testValue",
24+ },
25+ {
26+ name: "returns default when key not found",
27+ setup: func(ctx context.Context) context.Context {
28+ return ctx
29+ },
30+ key: "missingKey",
31+ defaultValue: "default",
32+ want: "default",
33+ },
34+ {
35+ name: "returns default when type mismatch",
36+ setup: func(ctx context.Context) context.Context {
37+ return context.WithValue(ctx, "testKey", 123) // int, not string
38+ },
39+ key: "testKey",
40+ defaultValue: "default",
41+ want: "default",
42+ },
43+ {
44+ name: "returns bool value",
45+ setup: func(ctx context.Context) context.Context {
46+ return context.WithValue(ctx, "boolKey", true)
47+ },
48+ key: "boolKey",
49+ defaultValue: false,
50+ want: true,
51+ },
52+ {
53+ name: "returns int value",
54+ setup: func(ctx context.Context) context.Context {
55+ return context.WithValue(ctx, "intKey", 42)
56+ },
57+ key: "intKey",
58+ defaultValue: 0,
59+ want: 42,
60+ },
61+ }
62+
63+ for _, tt := range tests {
64+ t.Run(tt.name, func(t *testing.T) {
65+ ctx := tt.setup(context.Background())
66+
67+ var got any
68+ switch tt.defaultValue.(type) {
69+ case string:
70+ got = getContextValue(ctx, tt.key, tt.defaultValue.(string))
71+ case bool:
72+ got = getContextValue(ctx, tt.key, tt.defaultValue.(bool))
73+ case int:
74+ got = getContextValue(ctx, tt.key, tt.defaultValue.(int))
75+ }
76+
77+ if got != tt.want {
78+ t.Errorf("getContextValue() = %v, want %v", got, tt.want)
79+ }
80+ })
81+ }
82+}
83+
84+func TestGetSessionFromContext(t *testing.T) {
85+ tests := []struct {
86+ name string
87+ ctx context.Context
88+ want string
89+ }{
90+ {
91+ name: "returns session ID when present",
92+ ctx: context.WithValue(context.Background(), SessionIDContextKey, "session-123"),
93+ want: "session-123",
94+ },
95+ {
96+ name: "returns empty string when not present",
97+ ctx: context.Background(),
98+ want: "",
99+ },
100+ {
101+ name: "returns empty string when wrong type",
102+ ctx: context.WithValue(context.Background(), SessionIDContextKey, 123),
103+ want: "",
104+ },
105+ }
106+
107+ for _, tt := range tests {
108+ t.Run(tt.name, func(t *testing.T) {
109+ got := GetSessionFromContext(tt.ctx)
110+ if got != tt.want {
111+ t.Errorf("GetSessionFromContext() = %v, want %v", got, tt.want)
112+ }
113+ })
114+ }
115+}
116+
117+func TestGetMessageFromContext(t *testing.T) {
118+ tests := []struct {
119+ name string
120+ ctx context.Context
121+ want string
122+ }{
123+ {
124+ name: "returns message ID when present",
125+ ctx: context.WithValue(context.Background(), MessageIDContextKey, "msg-456"),
126+ want: "msg-456",
127+ },
128+ {
129+ name: "returns empty string when not present",
130+ ctx: context.Background(),
131+ want: "",
132+ },
133+ {
134+ name: "returns empty string when wrong type",
135+ ctx: context.WithValue(context.Background(), MessageIDContextKey, 456),
136+ want: "",
137+ },
138+ }
139+
140+ for _, tt := range tests {
141+ t.Run(tt.name, func(t *testing.T) {
142+ got := GetMessageFromContext(tt.ctx)
143+ if got != tt.want {
144+ t.Errorf("GetMessageFromContext() = %v, want %v", got, tt.want)
145+ }
146+ })
147+ }
148+}
149+
150+func TestGetSupportsImagesFromContext(t *testing.T) {
151+ tests := []struct {
152+ name string
153+ ctx context.Context
154+ want bool
155+ }{
156+ {
157+ name: "returns true when present and true",
158+ ctx: context.WithValue(context.Background(), SupportsImagesContextKey, true),
159+ want: true,
160+ },
161+ {
162+ name: "returns false when present and false",
163+ ctx: context.WithValue(context.Background(), SupportsImagesContextKey, false),
164+ want: false,
165+ },
166+ {
167+ name: "returns false when not present",
168+ ctx: context.Background(),
169+ want: false,
170+ },
171+ {
172+ name: "returns false when wrong type",
173+ ctx: context.WithValue(context.Background(), SupportsImagesContextKey, "true"),
174+ want: false,
175+ },
176+ }
177+
178+ for _, tt := range tests {
179+ t.Run(tt.name, func(t *testing.T) {
180+ got := GetSupportsImagesFromContext(tt.ctx)
181+ if got != tt.want {
182+ t.Errorf("GetSupportsImagesFromContext() = %v, want %v", got, tt.want)
183+ }
184+ })
185+ }
186+}
187+
188+func TestGetModelNameFromContext(t *testing.T) {
189+ tests := []struct {
190+ name string
191+ ctx context.Context
192+ want string
193+ }{
194+ {
195+ name: "returns model name when present",
196+ ctx: context.WithValue(context.Background(), ModelNameContextKey, "claude-opus-4"),
197+ want: "claude-opus-4",
198+ },
199+ {
200+ name: "returns empty string when not present",
201+ ctx: context.Background(),
202+ want: "",
203+ },
204+ {
205+ name: "returns empty string when wrong type",
206+ ctx: context.WithValue(context.Background(), ModelNameContextKey, 789),
207+ want: "",
208+ },
209+ }
210+
211+ for _, tt := range tests {
212+ t.Run(tt.name, func(t *testing.T) {
213+ got := GetModelNameFromContext(tt.ctx)
214+ if got != tt.want {
215+ t.Errorf("GetModelNameFromContext() = %v, want %v", got, tt.want)
216+ }
217+ })
218+ }
219+}
internal/agent/tools/tools.go+17−35
@@ -22,53 +22,35 @@ const (
22 ModelNameContextKey modelNameKey = "model_name"
22 ModelNameContextKey modelNameKey = "model_name"
23)
23)
24
24
25−// GetSessionFromContext retrieves the session ID from the context.
25+// getContextValue is a generic helper that retrieves a typed value from context.
26−func GetSessionFromContext(ctx context.Context) string {
26+// If the value is not found or has the wrong type, it returns the default value.
27− sessionID := ctx.Value(SessionIDContextKey)
27+func getContextValue[T any](ctx context.Context, key any, defaultValue T) T {
28+ value := ctx.Value(key)
28− if sessionID == nil {
29+ if value == nil {
29− return ""
30+ return defaultValue
30 }
31 }
31− s, ok := sessionID.(string)
32+ if typedValue, ok := value.(T); ok {
32− if !ok {
33− return ""
33+ return typedValue
34 }
34 }
35− return s
35+ return defaultValue
36+}
37+
38+// GetSessionFromContext retrieves the session ID from the context.
39+func GetSessionFromContext(ctx context.Context) string {
40+ return getContextValue(ctx, SessionIDContextKey, "")
36}
41}
37
42
38// GetMessageFromContext retrieves the message ID from the context.
43// GetMessageFromContext retrieves the message ID from the context.
39func GetMessageFromContext(ctx context.Context) string {
44func GetMessageFromContext(ctx context.Context) string {
40− messageID := ctx.Value(MessageIDContextKey)
45+ return getContextValue(ctx, MessageIDContextKey, "")
41− if messageID == nil {
42− return ""
43− }
44− s, ok := messageID.(string)
45− if !ok {
46− return ""
47− }
48− return s
49}
46}
50
47
51// GetSupportsImagesFromContext retrieves whether the model supports images from the context.
48// GetSupportsImagesFromContext retrieves whether the model supports images from the context.
52func GetSupportsImagesFromContext(ctx context.Context) bool {
49func GetSupportsImagesFromContext(ctx context.Context) bool {
53− supportsImages := ctx.Value(SupportsImagesContextKey)
50+ return getContextValue(ctx, SupportsImagesContextKey, false)
54− if supportsImages == nil {
55− return false
56− }
57− if supports, ok := supportsImages.(bool); ok {
58− return supports
59− }
60− return false
61}
51}
62
52
63// GetModelNameFromContext retrieves the model name from the context.
53// GetModelNameFromContext retrieves the model name from the context.
64func GetModelNameFromContext(ctx context.Context) string {
54func GetModelNameFromContext(ctx context.Context) string {
65− modelName := ctx.Value(ModelNameContextKey)
55+ return getContextValue(ctx, ModelNameContextKey, "")
66− if modelName == nil {
67− return ""
68− }
69− s, ok := modelName.(string)
70− if !ok {
71− return ""
72− }
73− return s
74}
56}