95 lines
3.3 KiB
Go
95 lines
3.3 KiB
Go
package vikingmemory
|
|||
|
|
|
||
|
|
import (
|
||
|
|
"context"
|
||
|
|
"data-terminal-cloud/internal/outbox"
|
||
|
|
"encoding/json"
|
||
|
|
"fmt"
|
||
|
|
"net/http"
|
||
|
|
"net/http/httptest"
|
||
|
|
"regexp"
|
||
|
|
"strings"
|
||
|
|
"testing"
|
||
|
|
"time"
|
||
|
|
)
|
||
|
|
|
||
|
|
func testClient(t *testing.T, handler http.HandlerFunc) *Client {
|
||
|
|
t.Helper()
|
||
|
|
server := httptest.NewTLSServer(handler)
|
||
|
|
t.Cleanup(server.Close)
|
||
|
|
c, err := NewClient(Options{BaseURL: server.URL, ResourceID: "mem-test", APIKey: "test-secret", EventType: "event_v1", ProfileType: "profile_v1", Timeout: time.Second})
|
||
|
|
if err != nil {
|
||
|
|
t.Fatal(err)
|
||
|
|
}
|
||
|
|
c.http = server.Client()
|
||
|
|
return c
|
||
|
|
}
|
||
|
|
func TestSearchBindsUserAndRejectsForeignResults(t *testing.T) {
|
||
|
|
calls := 0
|
||
|
|
c := testClient(t, func(w http.ResponseWriter, r *http.Request) {
|
||
|
|
calls++
|
||
|
|
var input map[string]any
|
||
|
|
json.NewDecoder(r.Body).Decode(&input)
|
||
|
|
filter := input["filter"].(map[string]any)
|
||
|
|
if filter["user_id"].([]any)[0] != "alice" || filter["assistant_id"].([]any)[0] != AssistantID || input["resource_id"] != "mem-test" {
|
||
|
|
t.Errorf("incorrect scope: %v", input)
|
||
|
|
}
|
||
|
|
if r.Header.Get("Authorization") != "Bearer test-secret" {
|
||
|
|
t.Error("missing auth")
|
||
|
|
}
|
||
|
|
fmt.Fprint(w, `{"code":0,"data":{"result_list":[{"id":"foreign","user_id":["bob"],"assistant_id":["platform_companion_v1"],"memory_info":{"summary":"must not return"}},{"id":"ours","user_id":["alice"],"assistant_id":["platform_companion_v1"],"memory_info":{"summary":"喜欢茶"}}]}}`)
|
||
|
|
})
|
||
|
|
items, err := c.Search(context.Background(), "alice", "茶", 2)
|
||
|
|
if err != nil || len(items) != 2 || items[0].ID != "ours" || calls != 2 {
|
||
|
|
t.Fatalf("items=%v calls=%d err=%v", items, calls, err)
|
||
|
|
}
|
||
|
|
if _, err = c.Search(context.Background(), "", "", 1); err == nil {
|
||
|
|
t.Fatal("empty identity accepted")
|
||
|
|
}
|
||
|
|
}
|
||
|
|
func TestVikingErrorsDoNotLeakServerBodyOrKey(t *testing.T) {
|
||
|
|
for _, status := range []int{200, 401, 429, 503} {
|
||
|
|
t.Run(fmt.Sprint(status), func(t *testing.T) {
|
||
|
|
c := testClient(t, func(w http.ResponseWriter, r *http.Request) {
|
||
|
|
w.WriteHeader(status)
|
||
|
|
fmt.Fprint(w, `{"code":900,"message":"test-secret private transcript"}`)
|
||
|
|
})
|
||
|
|
err := c.AddSession(context.Background(), "c_test", "alice", []Message{{Role: "user", Content: "test"}}, 1)
|
||
|
|
if err == nil || strings.Contains(err.Error(), "secret") || strings.Contains(err.Error(), "transcript") {
|
||
|
|
t.Fatalf("unsafe error %v", err)
|
||
|
|
}
|
||
|
|
if outbox.IsRetryable(err) != (status != 401) {
|
||
|
|
t.Fatalf("retry classification %v", err)
|
||
|
|
}
|
||
|
|
})
|
||
|
|
}
|
||
|
|
}
|
||
|
|
func TestStableNamespacedIdentityAndSessionIDs(t *testing.T) {
|
||
|
|
if UserID("a", "x") == UserID("b", "x") || UserID("a", "x") == UserID("a", "y") {
|
||
|
|
t.Fatal("identity collision")
|
||
|
|
}
|
||
|
|
id := SessionID("12345678-1234-1234-1234-123456789abc", 1)
|
||
|
|
if !regexp.MustCompile(`^[a-zA-Z][a-zA-Z0-9_]{0,127}$`).MatchString(id) || id == SessionID("12345678-1234-1234-1234-123456789abc", 0) {
|
||
|
|
t.Fatal(id)
|
||
|
|
}
|
||
|
|
}
|
||
|
|
func TestDeleteUsesExactUserForBothMemoryTypes(t *testing.T) {
|
||
|
|
calls := 0
|
||
|
|
c := testClient(t, func(w http.ResponseWriter, r *http.Request) {
|
||
|
|
calls++
|
||
|
|
var input map[string]any
|
||
|
|
json.NewDecoder(r.Body).Decode(&input)
|
||
|
|
f := input["filter"].(map[string]any)
|
||
|
|
if f["user_id"].([]any)[0] != "alice" {
|
||
|
|
t.Error("unscoped delete")
|
||
|
|
}
|
||
|
|
fmt.Fprint(w, `{"code":0}`)
|
||
|
|
})
|
||
|
|
if err := c.DeleteUser(context.Background(), ""); err == nil || calls != 0 {
|
||
|
|
t.Fatal("empty deletion accepted")
|
||
|
|
}
|
||
|
|
if err := c.DeleteUser(context.Background(), "alice"); err != nil || calls != 2 {
|
||
|
|
t.Fatal(err, calls)
|
||
|
|
}
|
||
|
|
}
|