Skip to content

Commit bf92f70

Browse files
committed
fix tests
Signed-off-by: Yuval Kohavi <yuval.kohavi@gmail.com>
1 parent 0734274 commit bf92f70

File tree

6 files changed

+104
-17
lines changed

6 files changed

+104
-17
lines changed

go/internal/httpserver/auth/authn.go

Lines changed: 4 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -33,6 +33,9 @@ func AuthSessionFrom(ctx context.Context) (*Session, bool) {
3333
v, ok := ctx.Value(sessionKey).(*Session)
3434
return v, ok && v != nil
3535
}
36+
func AuthSessionTo(ctx context.Context, session *Session) context.Context {
37+
return context.WithValue(ctx, sessionKey, session)
38+
}
3639

3740
type Principal struct {
3841
User string
@@ -63,7 +66,7 @@ func AuthnMiddleware(authn AuthProvider) func(http.Handler) http.Handler {
6366
return
6467
}
6568
if session != nil {
66-
r = r.WithContext(context.WithValue(r.Context(), sessionKey, session))
69+
r = r.WithContext(AuthSessionTo(r.Context(), session))
6770
}
6871
next.ServeHTTP(w, r)
6972
})

go/internal/httpserver/handlers/agents_test.go

Lines changed: 15 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -18,6 +18,7 @@ import (
1818
"github.com/kagent-dev/kagent/go/internal/adk"
1919
"github.com/kagent-dev/kagent/go/internal/database"
2020
database_fake "github.com/kagent-dev/kagent/go/internal/database/fake"
21+
"github.com/kagent-dev/kagent/go/internal/httpserver/auth"
2122
"github.com/kagent-dev/kagent/go/internal/httpserver/handlers"
2223
common "github.com/kagent-dev/kagent/go/internal/utils"
2324
"github.com/kagent-dev/kagent/go/pkg/client/api"
@@ -66,7 +67,7 @@ func setupTestHandler(objects ...client.Object) (*handlers.AgentsHandler, string
6667
WithObjects(objects...).
6768
Build()
6869

69-
userID := common.GetGlobalUserID()
70+
userID := "test-user"
7071
dbClient := database_fake.NewClient()
7172

7273
base := &handlers.Base{
@@ -76,6 +77,7 @@ func setupTestHandler(objects ...client.Object) (*handlers.AgentsHandler, string
7677
Namespace: "default",
7778
},
7879
DatabaseService: dbClient,
80+
Authorizer: &auth.NothingAuthorizer{},
7981
}
8082

8183
return handlers.NewAgentsHandler(base), userID
@@ -99,6 +101,7 @@ func TestHandleGetAgent(t *testing.T) {
99101

100102
req := httptest.NewRequest("GET", "/api/agents/default/test-team", nil)
101103
req = mux.SetURLVars(req, map[string]string{"namespace": "default", "name": "test-team"})
104+
req = setUser(req, "test-user")
102105
w := httptest.NewRecorder()
103106

104107
handler.HandleGetAgent(&testErrorResponseWriter{w}, req)
@@ -136,6 +139,7 @@ func TestHandleGetAgent(t *testing.T) {
136139

137140
req := httptest.NewRequest("GET", "/api/agents/default/test-agent-ready", nil)
138141
req = mux.SetURLVars(req, map[string]string{"namespace": "default", "name": "test-agent-ready"})
142+
req = setUser(req, "test-user")
139143
w := httptest.NewRecorder()
140144

141145
handler.HandleGetAgent(&testErrorResponseWriter{w}, req)
@@ -164,6 +168,7 @@ func TestHandleGetAgent(t *testing.T) {
164168

165169
req := httptest.NewRequest("GET", "/api/agents/default/test-agent-not-ready", nil)
166170
req = mux.SetURLVars(req, map[string]string{"namespace": "default", "name": "test-agent-not-ready"})
171+
req = setUser(req, "test-user")
167172
w := httptest.NewRecorder()
168173

169174
handler.HandleGetAgent(&testErrorResponseWriter{w}, req)
@@ -192,6 +197,7 @@ func TestHandleGetAgent(t *testing.T) {
192197

193198
req := httptest.NewRequest("GET", "/api/agents/default/test-agent-different-reason", nil)
194199
req = mux.SetURLVars(req, map[string]string{"namespace": "default", "name": "test-agent-different-reason"})
200+
req = setUser(req, "test-user")
195201
w := httptest.NewRecorder()
196202

197203
handler.HandleGetAgent(&testErrorResponseWriter{w}, req)
@@ -209,6 +215,7 @@ func TestHandleGetAgent(t *testing.T) {
209215

210216
req := httptest.NewRequest("GET", "/api/agents/default/test-team", nil)
211217
req = mux.SetURLVars(req, map[string]string{"namespace": "default", "name": "test-team"})
218+
req = setUser(req, "test-user")
212219
w := httptest.NewRecorder()
213220

214221
handler.HandleGetAgent(&testErrorResponseWriter{w}, req)
@@ -239,6 +246,8 @@ func TestHandleListAgents(t *testing.T) {
239246
createAgent(handler.Base.DatabaseService, notReadyAgent)
240247

241248
req := httptest.NewRequest("GET", "/api/agents", nil)
249+
req = setUser(req, "test-user")
250+
242251
w := httptest.NewRecorder()
243252

244253
handler.HandleListAgents(&testErrorResponseWriter{w}, req)
@@ -290,6 +299,7 @@ func TestHandleUpdateAgent(t *testing.T) {
290299
req := httptest.NewRequest("PUT", "/api/agents/default/test-team", bytes.NewBuffer(body))
291300
req = mux.SetURLVars(req, map[string]string{"namespace": "default", "name": "test-team"})
292301
req.Header.Set("Content-Type", "application/json")
302+
req = setUser(req, "test-user")
293303
w := httptest.NewRecorder()
294304

295305
handler.HandleUpdateAgent(&testErrorResponseWriter{w}, req)
@@ -313,6 +323,7 @@ func TestHandleUpdateAgent(t *testing.T) {
313323
req := httptest.NewRequest("PUT", "/api/agents/default/non-existent", bytes.NewBuffer(body))
314324
req = mux.SetURLVars(req, map[string]string{"namespace": "default", "name": "non-existent"})
315325
req.Header.Set("Content-Type", "application/json")
326+
req = setUser(req, "test-user")
316327
w := httptest.NewRecorder()
317328

318329
handler.HandleUpdateAgent(&testErrorResponseWriter{w}, req)
@@ -349,6 +360,7 @@ func TestHandleCreateAgent(t *testing.T) {
349360
body, _ := json.Marshal(agent)
350361
req := httptest.NewRequest("POST", "/api/agents", bytes.NewBuffer(body))
351362
req.Header.Set("Content-Type", "application/json")
363+
req = setUser(req, "test-user")
352364
w := httptest.NewRecorder()
353365

354366
handler.HandleCreateAgent(&testErrorResponseWriter{w}, req)
@@ -376,6 +388,7 @@ func TestHandleDeleteTeam(t *testing.T) {
376388

377389
req := httptest.NewRequest("DELETE", "/api/agents/default/test-team", nil)
378390
req = mux.SetURLVars(req, map[string]string{"namespace": "default", "name": "test-team"})
391+
req = setUser(req, "test-user")
379392
w := httptest.NewRecorder()
380393

381394
handler.HandleDeleteAgent(&testErrorResponseWriter{w}, req)
@@ -391,6 +404,7 @@ func TestHandleDeleteTeam(t *testing.T) {
391404
"namespace": "default",
392405
"name": "non-existent",
393406
})
407+
req = setUser(req, "test-user")
394408
w := httptest.NewRecorder()
395409

396410
handler.HandleDeleteAgent(&testErrorResponseWriter{w}, req)

go/internal/httpserver/handlers/memory_test.go

Lines changed: 15 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -20,6 +20,7 @@ import (
2020

2121
"github.com/kagent-dev/kagent/go/controller/api/v1alpha1"
2222
database_fake "github.com/kagent-dev/kagent/go/internal/database/fake"
23+
"github.com/kagent-dev/kagent/go/internal/httpserver/auth"
2324
"github.com/kagent-dev/kagent/go/internal/httpserver/handlers"
2425
"github.com/kagent-dev/kagent/go/pkg/client/api"
2526
)
@@ -38,6 +39,7 @@ func TestMemoryHandler(t *testing.T) {
3839
KubeClient: kubeClient,
3940
DefaultModelConfig: types.NamespacedName{Namespace: "default", Name: "default"},
4041
DatabaseService: database_fake.NewClient(),
42+
Authorizer: &auth.NothingAuthorizer{},
4143
}
4244
handler := handlers.NewMemoryHandler(base)
4345
responseRecorder := newMockErrorResponseWriter()
@@ -72,6 +74,7 @@ func TestMemoryHandler(t *testing.T) {
7274
require.NoError(t, err)
7375

7476
req := httptest.NewRequest("GET", "/api/memories/", nil)
77+
req = setUser(req, "test-user")
7578
handler.HandleListMemories(responseRecorder, req)
7679

7780
assert.Equal(t, http.StatusOK, responseRecorder.Code)
@@ -93,6 +96,7 @@ func TestMemoryHandler(t *testing.T) {
9396
handler, _, responseRecorder := setupHandler()
9497

9598
req := httptest.NewRequest("GET", "/api/memories/", nil)
99+
req = setUser(req, "test-user")
96100
handler.HandleListMemories(responseRecorder, req)
97101

98102
assert.Equal(t, http.StatusOK, responseRecorder.Code)
@@ -123,6 +127,7 @@ func TestMemoryHandler(t *testing.T) {
123127

124128
jsonBody, _ := json.Marshal(reqBody)
125129
req := httptest.NewRequest("POST", "/api/memories/", bytes.NewBuffer(jsonBody))
130+
req = setUser(req, "test-user")
126131
req.Header.Set("Content-Type", "application/json")
127132

128133
handler.HandleCreateMemory(responseRecorder, req)
@@ -141,6 +146,7 @@ func TestMemoryHandler(t *testing.T) {
141146
handler, _, responseRecorder := setupHandler()
142147

143148
req := httptest.NewRequest("POST", "/api/memories/", bytes.NewBufferString("invalid json"))
149+
req = setUser(req, "test-user")
144150
req.Header.Set("Content-Type", "application/json")
145151

146152
handler.HandleCreateMemory(responseRecorder, req)
@@ -160,6 +166,7 @@ func TestMemoryHandler(t *testing.T) {
160166

161167
jsonBody, _ := json.Marshal(reqBody)
162168
req := httptest.NewRequest("POST", "/api/memories/", bytes.NewBuffer(jsonBody))
169+
req = setUser(req, "test-user")
163170
req.Header.Set("Content-Type", "application/json")
164171

165172
handler.HandleCreateMemory(responseRecorder, req)
@@ -192,6 +199,7 @@ func TestMemoryHandler(t *testing.T) {
192199

193200
jsonBody, _ := json.Marshal(reqBody)
194201
req := httptest.NewRequest("POST", "/api/memories/", bytes.NewBuffer(jsonBody))
202+
req = setUser(req, "test-user")
195203
req.Header.Set("Content-Type", "application/json")
196204

197205
handler.HandleCreateMemory(responseRecorder, req)
@@ -226,6 +234,7 @@ func TestMemoryHandler(t *testing.T) {
226234
require.NoError(t, err)
227235

228236
req := httptest.NewRequest("GET", "/api/memories/default/test-memory", nil)
237+
req = setUser(req, "test-user")
229238

230239
router := mux.NewRouter()
231240
router.HandleFunc("/api/memories/{namespace}/{name}", func(w http.ResponseWriter, r *http.Request) {
@@ -247,6 +256,7 @@ func TestMemoryHandler(t *testing.T) {
247256
handler, _, responseRecorder := setupHandler()
248257

249258
req := httptest.NewRequest("GET", "/api/memories/default/nonexistent", nil)
259+
req = setUser(req, "test-user")
250260

251261
router := mux.NewRouter()
252262
router.HandleFunc("/api/memories/{namespace}/{name}", func(w http.ResponseWriter, r *http.Request) {
@@ -291,6 +301,7 @@ func TestMemoryHandler(t *testing.T) {
291301

292302
jsonBody, _ := json.Marshal(reqBody)
293303
req := httptest.NewRequest("PUT", "/api/memories/default/test-memory", bytes.NewBuffer(jsonBody))
304+
req = setUser(req, "test-user")
294305
req.Header.Set("Content-Type", "application/json")
295306

296307
router := mux.NewRouter()
@@ -313,6 +324,7 @@ func TestMemoryHandler(t *testing.T) {
313324
handler, _, responseRecorder := setupHandler()
314325

315326
req := httptest.NewRequest("PUT", "/api/memories/default/test-memory", bytes.NewBufferString("invalid json"))
327+
req = setUser(req, "test-user")
316328
req.Header.Set("Content-Type", "application/json")
317329

318330
router := mux.NewRouter()
@@ -337,6 +349,7 @@ func TestMemoryHandler(t *testing.T) {
337349

338350
jsonBody, _ := json.Marshal(reqBody)
339351
req := httptest.NewRequest("PUT", "/api/memories/default/nonexistent", bytes.NewBuffer(jsonBody))
352+
req = setUser(req, "test-user")
340353
req.Header.Set("Content-Type", "application/json")
341354

342355
router := mux.NewRouter()
@@ -370,6 +383,7 @@ func TestMemoryHandler(t *testing.T) {
370383
require.NoError(t, err)
371384

372385
req := httptest.NewRequest("DELETE", "/api/memories/default/test-memory", nil)
386+
req = setUser(req, "test-user")
373387

374388
router := mux.NewRouter()
375389
router.HandleFunc("/api/memories/{namespace}/{name}", func(w http.ResponseWriter, r *http.Request) {
@@ -390,6 +404,7 @@ func TestMemoryHandler(t *testing.T) {
390404
handler, _, responseRecorder := setupHandler()
391405

392406
req := httptest.NewRequest("DELETE", "/api/memories/default/nonexistent", nil)
407+
req = setUser(req, "test-user")
393408

394409
router := mux.NewRouter()
395410
router.HandleFunc("/api/memories/{namespace}/{name}", func(w http.ResponseWriter, r *http.Request) {

0 commit comments

Comments
 (0)