Spaces:
Sleeping
Sleeping
Download src/internal/server/server_test.go from 9mark9/lethe-agent-memory: direct link, hf CLI and curl.
- Browser
- Download file 11.1 kB
-
https://huggingface.co/spaces/9mark9/lethe-agent-memory/resolve/main/src/internal/server/server_test.go
- Command line
-
hf download hf://spaces/9mark9/lethe-agent-memory/src/internal/server/server_test.go
-
curl -L -o server_test.go https://huggingface.co/spaces/9mark9/lethe-agent-memory/resolve/main/src/internal/server/server_test.go
11.1 kB
| package server | |
| import ( | |
| "encoding/json" | |
| "io" | |
| "log/slog" | |
| "net/http" | |
| "net/http/httptest" | |
| "os" | |
| "path/filepath" | |
| "strings" | |
| "sync" | |
| "testing" | |
| "time" | |
| "github.com/shaikn6/lethe/internal/memory" | |
| ) | |
| const adminToken = "admin-secret" | |
| var asAdmin = []string{"Authorization", "Bearer " + adminToken} | |
| func newHandler(t testing.TB, token string) http.Handler { | |
| return newHandlerWith(t, token, memory.Options{Seed: 1}) | |
| } | |
| func newHandlerWith(t testing.TB, token string, opts memory.Options) http.Handler { | |
| t.Helper() | |
| store, err := memory.Open(opts) | |
| if err != nil { | |
| t.Fatal(err) | |
| } | |
| t.Cleanup(func() { store.Close() }) | |
| cfg := Config{Token: token, AdminToken: adminToken, CompactEvery: time.Hour} | |
| return New(store, cfg, slog.New(slog.NewTextHandler(io.Discard, nil))) | |
| } | |
| func do(h http.Handler, method, path, body string, headers ...string) *httptest.ResponseRecorder { | |
| req := httptest.NewRequest(method, path, strings.NewReader(body)) | |
| for i := 0; i+1 < len(headers); i += 2 { | |
| req.Header.Set(headers[i], headers[i+1]) | |
| } | |
| rec := httptest.NewRecorder() | |
| h.ServeHTTP(rec, req) | |
| return rec | |
| } | |
| func TestRememberRecallFeedbackForget(t *testing.T) { | |
| h := newHandler(t, "") | |
| rec := do(h, "POST", "/v1/memories", `{"namespace":"u1","text":"Sam prefers window seats","embedding":[1,0],"metadata":{"source":"chat"}}`) | |
| if rec.Code != http.StatusCreated { | |
| t.Fatalf("remember: %d %s", rec.Code, rec.Body) | |
| } | |
| var stored struct { | |
| ID string `json:"id"` | |
| Importance float64 `json:"importance"` | |
| } | |
| json.Unmarshal(rec.Body.Bytes(), &stored) | |
| if stored.ID == "" || stored.Importance != 0.5 || strings.Contains(rec.Body.String(), "embedding") { | |
| t.Fatalf("remember body = %s", rec.Body) | |
| } | |
| rec = do(h, "POST", "/v1/recall", `{"namespace":"u1","query":"which seat does Sam like","embedding":[1,0],"explore":true}`) | |
| var got struct { | |
| RecallID string `json:"recall_id"` | |
| Memories []struct { | |
| ID string `json:"id"` | |
| Text string `json:"text"` | |
| } `json:"memories"` | |
| } | |
| json.Unmarshal(rec.Body.Bytes(), &got) | |
| if rec.Code != http.StatusOK || len(got.Memories) != 1 || got.Memories[0].ID != stored.ID { | |
| t.Fatalf("recall: %d %s", rec.Code, rec.Body) | |
| } | |
| if rec = do(h, "POST", "/v1/feedback", `{"recall_id":"`+got.RecallID+`","reward":1}`); rec.Code != http.StatusNoContent { | |
| t.Fatalf("feedback: %d %s", rec.Code, rec.Body) | |
| } | |
| if rec = do(h, "POST", "/v1/feedback", `{"recall_id":"`+got.RecallID+`","reward":1}`); rec.Code != http.StatusNotFound { | |
| t.Fatalf("repeated feedback: %d, want 404", rec.Code) | |
| } | |
| if rec = do(h, "GET", "/healthz", ""); !strings.Contains(rec.Body.String(), `"policy_updates":1`) || !strings.Contains(rec.Body.String(), `"memories":1`) { | |
| t.Fatalf("healthz = %s", rec.Body) | |
| } | |
| if rec = do(h, "DELETE", "/v1/memories/u1/"+stored.ID, ""); rec.Code != http.StatusNoContent { | |
| t.Fatalf("forget: %d %s", rec.Code, rec.Body) | |
| } | |
| if rec = do(h, "DELETE", "/v1/memories/u1/"+stored.ID, ""); rec.Code != http.StatusNotFound { | |
| t.Fatalf("second forget: %d, want 404", rec.Code) | |
| } | |
| rec = do(h, "POST", "/v1/recall", `{"namespace":"u1","query":"seat"}`) | |
| if !strings.Contains(rec.Body.String(), `"memories":[]`) { | |
| t.Fatalf("recall after forget = %s", rec.Body) | |
| } | |
| } | |
| func TestBadRequests(t *testing.T) { | |
| h := newHandler(t, "") | |
| cases := map[string]struct{ path, body string }{ | |
| "malformed json": {"/v1/memories", `{"namespace":`}, | |
| "unknown field": {"/v1/memories", `{"namespace":"n","text":"x","ttl":5}`}, | |
| "missing namespace": {"/v1/memories", `{"text":"x"}`}, | |
| "empty text": {"/v1/memories", `{"namespace":"n","text":""}`}, | |
| "importance range": {"/v1/memories", `{"namespace":"n","text":"x","importance":7}`}, | |
| "k out of range": {"/v1/recall", `{"namespace":"n","query":"x","k":5000}`}, | |
| "negative k": {"/v1/recall", `{"namespace":"n","query":"x","k":-2}`}, | |
| "feedback no reward": {"/v1/feedback", `{"recall_id":"abc"}`}, | |
| "reward out of range": {"/v1/feedback", `{"recall_id":"abc","reward":3}`}, | |
| "body too large": {"/v1/memories", `{"namespace":"n","text":"` + strings.Repeat("a", maxBodyBytes) + `"}`}, | |
| } | |
| for name, tc := range cases { | |
| if rec := do(h, "POST", tc.path, tc.body); rec.Code != http.StatusBadRequest { | |
| t.Errorf("%s: status = %d, want 400", name, rec.Code) | |
| } | |
| } | |
| if rec := do(h, "GET", "/v1/recall", ""); rec.Code != http.StatusMethodNotAllowed { | |
| t.Errorf("GET /v1/recall = %d, want 405", rec.Code) | |
| } | |
| } | |
| func TestBearerToken(t *testing.T) { | |
| h := newHandler(t, "s3cret") | |
| body := `{"namespace":"n","text":"x"}` | |
| if rec := do(h, "POST", "/v1/memories", body); rec.Code != http.StatusUnauthorized { | |
| t.Fatalf("no token: %d, want 401", rec.Code) | |
| } | |
| if rec := do(h, "POST", "/v1/memories", body, "Authorization", "Bearer wrong"); rec.Code != http.StatusUnauthorized { | |
| t.Fatalf("wrong token: %d, want 401", rec.Code) | |
| } | |
| if rec := do(h, "POST", "/v1/memories", body, "Authorization", "Bearer s3cret"); rec.Code != http.StatusCreated { | |
| t.Fatalf("right token: %d, want 201", rec.Code) | |
| } | |
| if rec := do(h, "GET", "/healthz", ""); rec.Code != http.StatusOK { | |
| t.Fatalf("healthz without token: %d, want 200", rec.Code) | |
| } | |
| } | |
| func TestPanicInHandlerBecomesA500(t *testing.T) { | |
| h := recoverPanics(slog.New(slog.NewTextHandler(io.Discard, nil)), http.HandlerFunc(func(http.ResponseWriter, *http.Request) { | |
| panic("boom") | |
| })) | |
| if rec := do(h, "GET", "/", ""); rec.Code != http.StatusInternalServerError { | |
| t.Fatalf("status = %d, want 500", rec.Code) | |
| } | |
| } | |
| // Arbitrary request bodies must never crash the server or produce a 500. | |
| func FuzzAPI(f *testing.F) { | |
| h := newHandler(f, "") | |
| do(h, "POST", "/v1/memories", `{"namespace":"n","text":"seed memory","embedding":[1,0]}`) | |
| for _, seed := range []string{ | |
| `{"namespace":"n","text":"hello","embedding":[0.5,0.5]}`, | |
| `{"namespace":"n","query":"hello","k":3,"explore":true}`, | |
| `{"recall_id":"x","reward":0.5}`, | |
| `{"namespace":"n","query":"","embedding":[1e308,1e308]}`, | |
| `{"namespace":"n","text":"x","created_at":"0001-01-01T00:00:00Z"}`, | |
| `[]`, `null`, `{"k":-1}`, ``, | |
| } { | |
| f.Add(seed) | |
| } | |
| f.Fuzz(func(t *testing.T, body string) { | |
| for _, path := range []string{"/v1/memories", "/v1/recall", "/v1/feedback"} { | |
| if rec := do(h, "POST", path, body); rec.Code >= 500 { | |
| t.Fatalf("POST %s with %q returned %d: %s", path, body, rec.Code, rec.Body) | |
| } | |
| } | |
| }) | |
| } | |
| func TestSetPolicyWeights(t *testing.T) { | |
| h := newHandler(t, "") | |
| if rec := do(h, "PUT", "/v1/policy", `{"weights":[1,0,3,0,0,0,0,0,0]}`, asAdmin...); rec.Code != http.StatusNoContent { | |
| t.Fatalf("set weights: %d %s", rec.Code, rec.Body) | |
| } | |
| if rec := do(h, "GET", "/healthz", ""); !strings.Contains(rec.Body.String(), `"weights":[1,0,3,0,0,0,0,0,0]`) { | |
| t.Fatalf("healthz = %s", rec.Body) | |
| } | |
| for name, body := range map[string]string{ | |
| "missing": `{}`, | |
| "too many": `{"weights":[1,2,3,4,5,6,7,8,9,10]}`, | |
| "too few": `{"weights":[1,2,3]}`, | |
| "too large": `{"weights":[1e9,0,0,0,0,0,0,0,0]}`, | |
| } { | |
| if rec := do(h, "PUT", "/v1/policy", body, asAdmin...); rec.Code != http.StatusBadRequest { | |
| t.Errorf("%s: status = %d, want 400", name, rec.Code) | |
| } | |
| } | |
| } | |
| func TestFullStoreReturns507(t *testing.T) { | |
| h := newHandlerWith(t, "", memory.Options{MaxMemories: 1}) | |
| body := `{"namespace":"n","text":"x"}` | |
| if rec := do(h, "POST", "/v1/memories", body); rec.Code != http.StatusCreated { | |
| t.Fatalf("first write: %d", rec.Code) | |
| } | |
| if rec := do(h, "POST", "/v1/memories", body); rec.Code != http.StatusInsufficientStorage { | |
| t.Fatalf("write to a full store: %d, want 507", rec.Code) | |
| } | |
| } | |
| func TestCompactErasesForgottenTextFromTheJournal(t *testing.T) { | |
| path := filepath.Join(t.TempDir(), "api.journal") | |
| h := newHandlerWith(t, "", memory.Options{JournalPath: path}) | |
| rec := do(h, "POST", "/v1/memories", `{"namespace":"n","text":"secret aardvark7 token"}`) | |
| var stored struct { | |
| ID string `json:"id"` | |
| } | |
| json.Unmarshal(rec.Body.Bytes(), &stored) | |
| do(h, "DELETE", "/v1/memories/n/"+stored.ID, "") | |
| if rec := do(h, "POST", "/v1/compact", "", asAdmin...); rec.Code != http.StatusNoContent { | |
| t.Fatalf("compact: %d %s", rec.Code, rec.Body) | |
| } | |
| if rec := do(h, "POST", "/v1/compact", "", asAdmin...); rec.Code != http.StatusTooManyRequests || rec.Header().Get("Retry-After") == "" { | |
| t.Fatalf("second compact right away: %d, want 429 with Retry-After", rec.Code) | |
| } | |
| data, err := os.ReadFile(path) | |
| if err != nil { | |
| t.Fatal(err) | |
| } | |
| if strings.Contains(string(data), "aardvark7") { | |
| t.Fatal("forgotten text is still on disk after compaction") | |
| } | |
| if rec := do(h, "GET", "/healthz", ""); rec.Code != http.StatusOK || !strings.Contains(rec.Body.String(), `"journal_healthy":true`) { | |
| t.Fatalf("healthz = %d %s", rec.Code, rec.Body) | |
| } | |
| } | |
| func TestAdminEndpointsNeedTheAdminToken(t *testing.T) { | |
| h := newHandler(t, "data-token") | |
| weights := `{"weights":[1,0,0,0,0,0,0,0,0]}` | |
| for name, headers := range map[string][]string{ | |
| "no token": nil, | |
| "data token": {"Authorization", "Bearer data-token"}, | |
| } { | |
| if rec := do(h, "PUT", "/v1/policy", weights, headers...); rec.Code != http.StatusUnauthorized { | |
| t.Errorf("PUT /v1/policy with %s: %d, want 401", name, rec.Code) | |
| } | |
| if rec := do(h, "POST", "/v1/compact", "", headers...); rec.Code != http.StatusUnauthorized { | |
| t.Errorf("POST /v1/compact with %s: %d, want 401", name, rec.Code) | |
| } | |
| } | |
| if rec := do(h, "POST", "/v1/memories", `{"namespace":"n","text":"x"}`, asAdmin...); rec.Code != http.StatusUnauthorized { | |
| t.Errorf("admin token on a data endpoint: %d, want 401", rec.Code) | |
| } | |
| } | |
| func TestAdminEndpointsAreOffWithoutAnAdminToken(t *testing.T) { | |
| store, err := memory.Open(memory.Options{}) | |
| if err != nil { | |
| t.Fatal(err) | |
| } | |
| h := New(store, Config{}, slog.New(slog.NewTextHandler(io.Discard, nil))) | |
| for _, req := range [][2]string{{"PUT", "/v1/policy"}, {"POST", "/v1/compact"}} { | |
| if rec := do(h, req[0], req[1], `{"weights":[1,0,0,0,0,0,0,0,0]}`); rec.Code != http.StatusNotFound && rec.Code != http.StatusMethodNotAllowed { | |
| t.Errorf("%s %s without an admin token configured: %d, want it unavailable", req[0], req[1], rec.Code) | |
| } | |
| } | |
| if rec := do(h, "GET", "/healthz", ""); !strings.Contains(rec.Body.String(), `"weights":[2,2,0,0,0,0,0,0,0]`) { | |
| t.Fatalf("weights changed without an admin token: %s", rec.Body) | |
| } | |
| } | |
| func TestSheds503WhenAtCapacity(t *testing.T) { | |
| release, entered := make(chan struct{}), make(chan struct{}) | |
| h := limitInFlight(1, http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { | |
| if r.URL.Path == "/slow" { | |
| close(entered) | |
| <-release | |
| } | |
| })) | |
| var wg sync.WaitGroup | |
| wg.Add(1) | |
| go func() { | |
| defer wg.Done() | |
| do(h, "GET", "/slow", "") | |
| }() | |
| <-entered | |
| if rec := do(h, "GET", "/v1/recall", ""); rec.Code != http.StatusServiceUnavailable || rec.Header().Get("Retry-After") == "" { | |
| t.Errorf("request beyond capacity: %d, want 503 with Retry-After", rec.Code) | |
| } | |
| if rec := do(h, "GET", "/healthz", ""); rec.Code != http.StatusOK { | |
| t.Errorf("health check at capacity: %d, want 200", rec.Code) | |
| } | |
| close(release) | |
| wg.Wait() | |
| if rec := do(h, "GET", "/v1/recall", ""); rec.Code != http.StatusOK { | |
| t.Errorf("after the slow request finished: %d, want 200", rec.Code) | |
| } | |
| } | |