Cyrene_v2/internal/runtime/persist/persist_test.go

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)
	}
}