feat(persist): durable checkpoint（SQLite upsert/LoadLatest/List/Delete，ErrNotFound）
This commit is contained in:
Aska committed
1 parent
ecd3cdf28c
commit
011d25a53d
2 files changed
+652
No files matched your search
@@ -0,0 +1,362 @@
|
||||
package persist
|
||||
|
||||
import (
|
||||
"context"
|
||||
"errors"
|
||||
"fmt"
|
||||
"io"
|
||||
"log/slog"
|
||||
"sync"
|
||||
"testing"
|
||||
"time"
|
||||
|
||||
dbpkg "git.yeij.top/AskaEth/Cyrene_v2/internal/db"
|
||||
)
|
||||
|
||||
// testLogger returns a logger that writes nowhere, keeping test output clean.
|
||||
func testLogger() *slog.Logger {
|
||||
return slog.New(slog.NewTextHandler(io.Discard, nil))
|
||||
}
|
||||
|
||||
// newTestStore opens an isolated in-memory SQLite database and returns a Store
|
||||
// backed by it. Each call owns a private database, so parallel tests do not
|
||||
// interfere with one another.
|
||||
func newTestStore(t *testing.T) *Store {
|
||||
t.Helper()
|
||||
ctx := context.Background()
|
||||
|
||||
db, err := dbpkg.Open(ctx, ":memory:", testLogger())
|
||||
if err != nil {
|
||||
t.Fatalf("open in-memory db: %v", err)
|
||||
}
|
||||
t.Cleanup(func() { _ = db.Close() })
|
||||
|
||||
s, err := New(ctx, db, testLogger())
|
||||
if err != nil {
|
||||
t.Fatalf("New(store) error: %v", err)
|
||||
}
|
||||
return s
|
||||
}
|
||||
|
||||
func TestStore_SaveLoad_正常(t *testing.T) {
|
||||
t.Parallel()
|
||||
ctx := context.Background()
|
||||
s := newTestStore(t)
|
||||
|
||||
cp := &Checkpoint{
|
||||
ID: "cp-1",
|
||||
AgentID: "agent-a",
|
||||
Step: 3,
|
||||
State: []byte(`{"turn":3,"msg":"hello"}`),
|
||||
}
|
||||
if err := s.Save(ctx, cp); err != nil {
|
||||
t.Fatalf("Save error: %v", err)
|
||||
}
|
||||
|
||||
got, err := s.Load(ctx, "cp-1")
|
||||
if err != nil {
|
||||
t.Fatalf("Load error: %v", err)
|
||||
}
|
||||
if got.ID != cp.ID || got.AgentID != cp.AgentID || got.Step != cp.Step {
|
||||
t.Fatalf("round-trip mismatch: got %+v", got)
|
||||
}
|
||||
if string(got.State) != string(cp.State) {
|
||||
t.Fatalf("state mismatch: got %q want %q", got.State, cp.State)
|
||||
}
|
||||
if got.CreatedAt.IsZero() || got.UpdatedAt.IsZero() {
|
||||
t.Fatalf("timestamps not set: created=%v updated=%v", got.CreatedAt, got.UpdatedAt)
|
||||
}
|
||||
}
|
||||
|
||||
func TestStore_Load_Missing_异常(t *testing.T) {
|
||||
t.Parallel()
|
||||
ctx := context.Background()
|
||||
s := newTestStore(t)
|
||||
|
||||
if _, err := s.Load(ctx, "does-not-exist"); !errors.Is(err, ErrNotFound) {
|
||||
t.Fatalf("Load(missing) err = %v, want wrap of ErrNotFound", err)
|
||||
}
|
||||
}
|
||||
|
||||
func TestStore_LoadLatest_正常(t *testing.T) {
|
||||
t.Parallel()
|
||||
ctx := context.Background()
|
||||
s := newTestStore(t)
|
||||
|
||||
// Save steps in ascending order; the latest save must be the newest step.
|
||||
for step := 1; step <= 3; step++ {
|
||||
cp := &Checkpoint{
|
||||
ID: fmt.Sprintf("cp-%d", step),
|
||||
AgentID: "agent-a",
|
||||
Step: step,
|
||||
State: []byte(fmt.Sprintf(`{"step":%d}`, step)),
|
||||
}
|
||||
if err := s.Save(ctx, cp); err != nil {
|
||||
t.Fatalf("Save(step %d) error: %v", step, err)
|
||||
}
|
||||
}
|
||||
|
||||
got, err := s.LoadLatest(ctx, "agent-a")
|
||||
if err != nil {
|
||||
t.Fatalf("LoadLatest error: %v", err)
|
||||
}
|
||||
if got.Step != 3 {
|
||||
t.Fatalf("LoadLatest step = %d, want 3", got.Step)
|
||||
}
|
||||
if got.ID != "cp-3" {
|
||||
t.Fatalf("LoadLatest id = %q, want cp-3", got.ID)
|
||||
}
|
||||
}
|
||||
|
||||
func TestStore_LoadLatest_Missing_异常(t *testing.T) {
|
||||
t.Parallel()
|
||||
ctx := context.Background()
|
||||
s := newTestStore(t)
|
||||
|
||||
if _, err := s.LoadLatest(ctx, "no-such-agent"); !errors.Is(err, ErrNotFound) {
|
||||
t.Fatalf("LoadLatest(missing) err = %v, want wrap of ErrNotFound", err)
|
||||
}
|
||||
}
|
||||
|
||||
func TestStore_List_正常(t *testing.T) {
|
||||
t.Parallel()
|
||||
ctx := context.Background()
|
||||
s := newTestStore(t)
|
||||
|
||||
// agent-a gets steps 1 and 2; agent-b gets a single step.
|
||||
for _, cp := range []*Checkpoint{
|
||||
{ID: "a-1", AgentID: "agent-a", Step: 1, State: []byte(`{"step":1}`)},
|
||||
{ID: "a-2", AgentID: "agent-a", Step: 2, State: []byte(`{"step":2}`)},
|
||||
{ID: "b-1", AgentID: "agent-b", Step: 1, State: []byte(`{"step":1}`)},
|
||||
} {
|
||||
if err := s.Save(ctx, cp); err != nil {
|
||||
t.Fatalf("Save(%s) error: %v", cp.ID, err)
|
||||
}
|
||||
}
|
||||
|
||||
listA, err := s.List(ctx, "agent-a")
|
||||
if err != nil {
|
||||
t.Fatalf("List(agent-a) error: %v", err)
|
||||
}
|
||||
if len(listA) != 2 {
|
||||
t.Fatalf("List(agent-a) length = %d, want 2", len(listA))
|
||||
}
|
||||
if listA[0].Step != 1 || listA[1].Step != 2 {
|
||||
t.Fatalf("List(agent-a) not ascending by step: %+v", listA)
|
||||
}
|
||||
|
||||
listB, err := s.List(ctx, "agent-b")
|
||||
if err != nil {
|
||||
t.Fatalf("List(agent-b) error: %v", err)
|
||||
}
|
||||
if len(listB) != 1 {
|
||||
t.Fatalf("List(agent-b) length = %d, want 1", len(listB))
|
||||
}
|
||||
|
||||
// An agent with no checkpoints returns an empty, non-nil slice.
|
||||
none, err := s.List(ctx, "no-such-agent")
|
||||
if err != nil {
|
||||
t.Fatalf("List(none) error: %v", err)
|
||||
}
|
||||
if none == nil {
|
||||
t.Fatalf("List(none) returned nil slice, want empty non-nil")
|
||||
}
|
||||
if len(none) != 0 {
|
||||
t.Fatalf("List(none) length = %d, want 0", len(none))
|
||||
}
|
||||
}
|
||||
|
||||
func TestStore_Save_Upsert_正常(t *testing.T) {
|
||||
t.Parallel()
|
||||
ctx := context.Background()
|
||||
s := newTestStore(t)
|
||||
|
||||
// A clearly past timestamp: the store must preserve it across the upsert
|
||||
// while refreshing UpdatedAt to the current time.
|
||||
originalCreated := time.Date(2024, time.January, 1, 0, 0, 0, 0, time.UTC)
|
||||
first := &Checkpoint{
|
||||
ID: "cp-1",
|
||||
AgentID: "agent-a",
|
||||
Step: 1,
|
||||
State: []byte(`{"step":1}`),
|
||||
CreatedAt: originalCreated,
|
||||
}
|
||||
if err := s.Save(ctx, first); err != nil {
|
||||
t.Fatalf("Save(first) error: %v", err)
|
||||
}
|
||||
|
||||
// A second save with the same ID must keep the original CreatedAt while
|
||||
// refreshing step, state and UpdatedAt.
|
||||
second := &Checkpoint{
|
||||
ID: "cp-1",
|
||||
AgentID: "agent-a",
|
||||
Step: 2,
|
||||
State: []byte(`{"step":2}`),
|
||||
}
|
||||
if err := s.Save(ctx, second); err != nil {
|
||||
t.Fatalf("Save(second) error: %v", err)
|
||||
}
|
||||
|
||||
got, err := s.Load(ctx, "cp-1")
|
||||
if err != nil {
|
||||
t.Fatalf("Load error: %v", err)
|
||||
}
|
||||
if got.Step != 2 {
|
||||
t.Fatalf("upsert step = %d, want 2", got.Step)
|
||||
}
|
||||
if string(got.State) != `{"step":2}` {
|
||||
t.Fatalf("upsert state = %q, want step 2", got.State)
|
||||
}
|
||||
if !got.CreatedAt.Equal(originalCreated) {
|
||||
t.Fatalf("upsert created_at = %v, want preserved %v", got.CreatedAt, originalCreated)
|
||||
}
|
||||
if got.UpdatedAt.Before(got.CreatedAt) {
|
||||
t.Fatalf("upsert updated_at %v is before created_at %v", got.UpdatedAt, got.CreatedAt)
|
||||
}
|
||||
|
||||
// Only one row must remain for the ID after the upsert.
|
||||
all, err := s.List(ctx, "agent-a")
|
||||
if err != nil {
|
||||
t.Fatalf("List(after upsert) error: %v", err)
|
||||
}
|
||||
if len(all) != 1 {
|
||||
t.Fatalf("upsert left %d rows, want 1", len(all))
|
||||
}
|
||||
}
|
||||
|
||||
func TestStore_Delete_正常(t *testing.T) {
|
||||
t.Parallel()
|
||||
ctx := context.Background()
|
||||
s := newTestStore(t)
|
||||
|
||||
cp := &Checkpoint{ID: "cp-1", AgentID: "agent-a", Step: 1, State: []byte(`{}`)}
|
||||
if err := s.Save(ctx, cp); err != nil {
|
||||
t.Fatalf("Save error: %v", err)
|
||||
}
|
||||
if err := s.Delete(ctx, "cp-1"); err != nil {
|
||||
t.Fatalf("Delete error: %v", err)
|
||||
}
|
||||
if _, err := s.Load(ctx, "cp-1"); !errors.Is(err, ErrNotFound) {
|
||||
t.Fatalf("Load after delete err = %v, want ErrNotFound", err)
|
||||
}
|
||||
}
|
||||
|
||||
func TestStore_Delete_Missing_异常(t *testing.T) {
|
||||
t.Parallel()
|
||||
ctx := context.Background()
|
||||
s := newTestStore(t)
|
||||
|
||||
if err := s.Delete(ctx, "never-existed"); !errors.Is(err, ErrNotFound) {
|
||||
t.Fatalf("Delete(missing) err = %v, want wrap of ErrNotFound", err)
|
||||
}
|
||||
}
|
||||
|
||||
func TestStore_Save_SQLInjection_安全(t *testing.T) {
|
||||
t.Parallel()
|
||||
ctx := context.Background()
|
||||
s := newTestStore(t)
|
||||
|
||||
// Malicious strings must be treated as opaque data, never executed as SQL.
|
||||
cp := &Checkpoint{
|
||||
ID: "1'; DROP TABLE checkpoints;--",
|
||||
AgentID: "agent'; DELETE FROM checkpoints;--",
|
||||
Step: 1,
|
||||
State: []byte(`{"x":"'); DROP TABLE checkpoints;--"}`),
|
||||
}
|
||||
if err := s.Save(ctx, cp); err != nil {
|
||||
t.Fatalf("Save(injection) error: %v", err)
|
||||
}
|
||||
|
||||
got, err := s.Load(ctx, cp.ID)
|
||||
if err != nil {
|
||||
t.Fatalf("Load(injection id) error: %v", err)
|
||||
}
|
||||
if got.ID != cp.ID || got.AgentID != cp.AgentID {
|
||||
t.Fatalf("injection round-trip mismatch: got %+v", got)
|
||||
}
|
||||
if string(got.State) != string(cp.State) {
|
||||
t.Fatalf("injection state mismatch: got %q", got.State)
|
||||
}
|
||||
|
||||
// The table must still exist and hold exactly the injected row.
|
||||
var tableCount int
|
||||
if err := s.db.QueryRowContext(ctx,
|
||||
"SELECT count(*) FROM sqlite_master WHERE type='table' AND name='checkpoints'").Scan(&tableCount); err != nil {
|
||||
t.Fatalf("query sqlite_master: %v", err)
|
||||
}
|
||||
if tableCount != 1 {
|
||||
t.Fatalf("checkpoints table missing after injection, count=%d", tableCount)
|
||||
}
|
||||
var rowCount int
|
||||
if err := s.db.QueryRowContext(ctx, "SELECT count(*) FROM checkpoints").Scan(&rowCount); err != nil {
|
||||
t.Fatalf("count checkpoints: %v", err)
|
||||
}
|
||||
if rowCount != 1 {
|
||||
t.Fatalf("checkpoints row count = %d, want 1", rowCount)
|
||||
}
|
||||
}
|
||||
|
||||
func TestStore_SaveConcurrent_并发(t *testing.T) {
|
||||
t.Parallel()
|
||||
ctx := context.Background()
|
||||
s := newTestStore(t)
|
||||
|
||||
const agents = 8
|
||||
const perAgent = 5
|
||||
|
||||
var wg sync.WaitGroup
|
||||
errCh := make(chan error, agents*perAgent)
|
||||
for a := 0; a < agents; a++ {
|
||||
agentID := fmt.Sprintf("agent-%d", a)
|
||||
for step := 0; step < perAgent; step++ {
|
||||
wg.Add(1)
|
||||
go func(id string, st int) {
|
||||
defer wg.Done()
|
||||
cp := &Checkpoint{
|
||||
ID: fmt.Sprintf("cp-%s-%d", id, st),
|
||||
AgentID: id,
|
||||
Step: st,
|
||||
State: []byte(fmt.Sprintf(`{"agent":%q,"step":%d}`, id, st)),
|
||||
}
|
||||
if err := s.Save(ctx, cp); err != nil {
|
||||
errCh <- err
|
||||
}
|
||||
}(agentID, step)
|
||||
}
|
||||
}
|
||||
wg.Wait()
|
||||
close(errCh)
|
||||
for err := range errCh {
|
||||
t.Fatalf("concurrent save failed: %v", err)
|
||||
}
|
||||
|
||||
for a := 0; a < agents; a++ {
|
||||
agentID := fmt.Sprintf("agent-%d", a)
|
||||
all, err := s.List(ctx, agentID)
|
||||
if err != nil {
|
||||
t.Fatalf("List(%s) error: %v", agentID, err)
|
||||
}
|
||||
if len(all) != perAgent {
|
||||
t.Fatalf("List(%s) length = %d, want %d", agentID, len(all), perAgent)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func TestStore_New_NilDB_异常(t *testing.T) {
|
||||
t.Parallel()
|
||||
ctx := context.Background()
|
||||
|
||||
if _, err := New(ctx, nil, testLogger()); !errors.Is(err, ErrNilDB) {
|
||||
t.Fatalf("New(nil db) err = %v, want wrap of ErrNilDB", err)
|
||||
}
|
||||
}
|
||||
|
||||
func TestStore_Save_MissingID_异常(t *testing.T) {
|
||||
t.Parallel()
|
||||
ctx := context.Background()
|
||||
s := newTestStore(t)
|
||||
|
||||
if err := s.Save(ctx, &Checkpoint{AgentID: "a", Step: 1, State: []byte(`{}`)}); !errors.Is(err, ErrMissingID) {
|
||||
t.Fatalf("Save(blank id) err = %v, want wrap of ErrMissingID", err)
|
||||
}
|
||||
}
|
||||
Reference in new issue
Block a user