Co-Authored-By: Claude Opus 5.5 (1M context) <noreply@anthropic.com>
78 lines
3.0 KiB
Go
78 lines
3.0 KiB
Go
package store
|
|
|
|
import (
|
|
"strings"
|
|
"testing"
|
|
"time"
|
|
)
|
|
|
|
func TestForgettingAPersonLeavesNothingThatNamesThem(t *testing.T) {
|
|
checkForget(t, capacityStore(t))
|
|
}
|
|
|
|
func TestPostgresForgettingAPersonLeavesNothingThatNamesThem(t *testing.T) {
|
|
checkForget(t, openPG(t))
|
|
}
|
|
|
|
func checkForget(t *testing.T, db *DB) {
|
|
ctx := t.Context()
|
|
now := time.Now()
|
|
org, _ := db.CreateOrg(ctx, "acme")
|
|
team, _ := db.CreateTeam(ctx, org, "ml")
|
|
project, _ := db.CreateProject(ctx, team, "train")
|
|
u := User{ID: "usr-ayse", Email: "ayse@example.com", State: UserActive, CreatedAt: now}
|
|
if err := db.CreateUser(ctx, u); err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
other := User{ID: "usr-bob", Email: "bob@example.com", State: UserActive, CreatedAt: now}
|
|
_ = db.CreateUser(ctx, other)
|
|
rows := []Usage{
|
|
{At: now, KeyID: "k1", UserID: u.ID, ProjectID: project, Model: "m", ServedModel: "m", PromptTokens: 10, CostMicros: 5, Status: 200, RequestID: "r-ayse"},
|
|
{At: now, KeyID: "session:" + u.ID, UserID: u.ID, ProjectID: project, Model: "m", ServedModel: "m", PromptTokens: 1, Status: 200},
|
|
{At: now, KeyID: "k2", UserID: other.ID, ProjectID: project, Model: "m", ServedModel: "m", PromptTokens: 7, Status: 200, RequestID: "r-bob"},
|
|
}
|
|
if err := db.WriteUsage(ctx, rows); err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
for _, id := range []string{"r-ayse", "r-bob"} {
|
|
_ = db.PutRequestBody(ctx, RequestBody{RequestID: id, ProjectID: project, At: now, ExpiresAt: now.Add(time.Hour),
|
|
Path: "/v1/chat/completions", Model: "m", ServedModel: "m", Status: 200, Request: "{}", Response: "x"})
|
|
}
|
|
_ = db.PutFile(ctx, File{ID: "file-a", Owner: "user:" + u.ID, Purpose: "batch", Filename: "f", CreatedAt: now}, []byte("x"))
|
|
_ = db.WriteAudit(ctx, AuditEvent{Actor: "admin", Source: "api", Action: "POST /api/v1/users", Target: u.Email,
|
|
Detail: `{"email":"ayse@example.com"}`, Status: 200})
|
|
|
|
f, err := db.ForgetUser(ctx, u.ID, now)
|
|
if err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
if f.UsageRows != 2 || f.Requests != 1 || f.Files != 1 || f.AuditEvents != 1 || !strings.HasPrefix(f.Pseudonym, "forgotten-") {
|
|
t.Errorf("%+v", f)
|
|
}
|
|
if _, err := db.UserByID(ctx, u.ID); err == nil {
|
|
t.Error("the account is still there")
|
|
}
|
|
var n int
|
|
for _, q := range []string{
|
|
`SELECT COUNT(*) FROM usage WHERE user_id = 'usr-ayse' OR key_id = 'session:usr-ayse'`,
|
|
`SELECT COUNT(*) FROM usage_1h WHERE user_id = 'usr-ayse' OR key_id = 'session:usr-ayse'`,
|
|
`SELECT COUNT(*) FROM usage_1d WHERE user_id = 'usr-ayse'`,
|
|
`SELECT COUNT(*) FROM audit_events WHERE target LIKE '%ayse%' OR detail LIKE '%ayse%'`,
|
|
`SELECT COUNT(*) FROM request_bodies WHERE request_id = 'r-ayse'`,
|
|
} {
|
|
_ = db.db.QueryRowContext(ctx, q).Scan(&n)
|
|
if n != 0 {
|
|
t.Errorf("%s: %d", q, n)
|
|
}
|
|
}
|
|
// The totals stay true, and someone else's are untouched.
|
|
var total int64
|
|
_ = db.db.QueryRowContext(ctx, `SELECT SUM(prompt_tokens) FROM usage WHERE user_id = ?`, f.Pseudonym).Scan(&total)
|
|
if total != 11 {
|
|
t.Errorf("their usage under the pseudonym: %d", total)
|
|
}
|
|
if _, err := db.RequestBodyByID(ctx, "r-bob"); err != nil {
|
|
t.Error("another person's kept request went too")
|
|
}
|
|
}
|