82 lines
3.7 KiB
Go
82 lines
3.7 KiB
Go
package retrieval
|
|
|
|
import (
|
|
"context"
|
|
"os"
|
|
"strings"
|
|
"testing"
|
|
|
|
"data-terminal-cloud/internal/developmentembedding"
|
|
|
|
"github.com/google/uuid"
|
|
"github.com/jackc/pgx/v5/pgxpool"
|
|
)
|
|
|
|
func TestSelectMemoryForContextAppliesEligibilityBoundary(t *testing.T) {
|
|
databaseURL := os.Getenv("TEST_DATABASE_URL")
|
|
if databaseURL == "" {
|
|
t.Skip("TEST_DATABASE_URL is not set")
|
|
}
|
|
ctx := context.Background()
|
|
pool, err := pgxpool.New(ctx, databaseURL)
|
|
if err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
defer pool.Close()
|
|
type fixture struct{ memoryID, memoryPublicID, versionID, versionPublicID, sensitivity, content string }
|
|
fixtures := []fixture{
|
|
{uuid.NewString(), uuid.NewString(), uuid.NewString(), uuid.NewString(), "normal", "长者在河边小学教书。"},
|
|
{uuid.NewString(), uuid.NewString(), uuid.NewString(), uuid.NewString(), "sensitive", "合成敏感记忆不得进入语音上下文。"},
|
|
}
|
|
defer func() {
|
|
for _, item := range fixtures {
|
|
_, _ = pool.Exec(ctx, `DELETE FROM memory_index_entry WHERE memory_id=$1::uuid`, item.memoryID)
|
|
_, _ = pool.Exec(ctx, `UPDATE curated_memory SET current_version_id=NULL WHERE id=$1::uuid`, item.memoryID)
|
|
_, _ = pool.Exec(ctx, `DELETE FROM memory_version WHERE memory_id=$1::uuid`, item.memoryID)
|
|
_, _ = pool.Exec(ctx, `DELETE FROM curated_memory WHERE id=$1::uuid`, item.memoryID)
|
|
}
|
|
}()
|
|
for _, item := range fixtures {
|
|
if _, err := pool.Exec(ctx, `
|
|
INSERT INTO curated_memory(id,public_id,tenant_id,subject_id,status,sensitivity,valid_from)
|
|
VALUES($1::uuid,$2::uuid,'10000000-0000-0000-0000-000000000001','10000000-0000-0000-0000-000000000401','ACTIVE',$3,now())`,
|
|
item.memoryID, item.memoryPublicID, item.sensitivity); err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
if _, err := pool.Exec(ctx, `
|
|
INSERT INTO memory_version(id,public_id,memory_id,version_no,content,source_refs,confirmed_by,content_digest)
|
|
VALUES($1::uuid,$2::uuid,$3::uuid,1,$4,'[{"type":"integration"}]','10000000-0000-0000-0000-000000000101',$5)`,
|
|
item.versionID, item.versionPublicID, item.memoryID, item.content, strings.Repeat("b", 64)); err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
if _, err := pool.Exec(ctx, `UPDATE curated_memory SET current_version_id=$1::uuid WHERE id=$2::uuid`, item.versionID, item.memoryID); err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
if _, err := pool.Exec(ctx, `
|
|
INSERT INTO memory_index_entry(tenant_id,subject_id,memory_id,memory_version_id,index_type,embedding,status,provider_ref)
|
|
VALUES('10000000-0000-0000-0000-000000000001','10000000-0000-0000-0000-000000000401',$1::uuid,$2::uuid,'pgvector',$3::vector,'ACTIVE','integration')`,
|
|
item.memoryID, item.versionID, developmentembedding.Literal(developmentembedding.Vector(item.content))); err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
}
|
|
tx, err := pool.Begin(ctx)
|
|
if err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
defer tx.Rollback(ctx)
|
|
if _, err := SelectMemoryForContext(ctx, tx, "10000000-0000-0000-0000-000000000001", "10000000-0000-0000-0000-000000000401", "河边小学", 0); err == nil {
|
|
t.Fatal("invalid retrieval limit was accepted")
|
|
}
|
|
items, err := SelectMemoryForContext(ctx, tx, "10000000-0000-0000-0000-000000000001", "10000000-0000-0000-0000-000000000401", "河边小学", 5)
|
|
if err != nil || len(items) != 1 {
|
|
t.Fatalf("eligible memories=%+v err=%v", items, err)
|
|
}
|
|
if items[0].MemoryID != fixtures[0].memoryPublicID || items[0].MemoryVersionID != fixtures[0].versionPublicID || items[0].Content != fixtures[0].content || items[0].Sensitivity != "normal" {
|
|
t.Fatalf("retrieved memory=%+v", items[0])
|
|
}
|
|
otherSubjectItems, err := SelectMemoryForContext(ctx, tx, "10000000-0000-0000-0000-000000000001", "10000000-0000-0000-0000-000000000402", "河边小学", 5)
|
|
if err != nil || len(otherSubjectItems) != 0 {
|
|
t.Fatalf("cross-subject memories=%+v err=%v", otherSubjectItems, err)
|
|
}
|
|
}
|