|
| 1 | +package storage |
| 2 | + |
| 3 | +import ( |
| 4 | + "fmt" |
| 5 | + "strings" |
| 6 | + "testing" |
| 7 | + "time" |
| 8 | + |
| 9 | + "github.com/stretchr/testify/assert" |
| 10 | + "github.com/stretchr/testify/require" |
| 11 | +) |
| 12 | + |
| 13 | +// saveSizedActivities saves n records (id-00 oldest → id-(n-1) newest), each |
| 14 | +// padded so its stored value is ~payload bytes. |
| 15 | +func saveSizedActivities(t *testing.T, m *Manager, n, payload int) { |
| 16 | + t.Helper() |
| 17 | + base := time.Date(2026, 1, 1, 0, 0, 0, 0, time.UTC) |
| 18 | + for i := 0; i < n; i++ { |
| 19 | + require.NoError(t, m.SaveActivity(&ActivityRecord{ |
| 20 | + ID: fmt.Sprintf("id-%02d", i), |
| 21 | + Type: ActivityTypeToolCall, |
| 22 | + Status: "success", |
| 23 | + Response: strings.Repeat("x", payload), |
| 24 | + Timestamp: base.Add(time.Duration(i) * time.Minute), |
| 25 | + })) |
| 26 | + } |
| 27 | +} |
| 28 | + |
| 29 | +// exists reports whether an activity record with the given id is present. |
| 30 | +// GetActivity returns (nil, nil) — not an error — when a record is absent. |
| 31 | +func exists(t *testing.T, m *Manager, id string) bool { |
| 32 | + t.Helper() |
| 33 | + rec, err := m.GetActivity(id) |
| 34 | + require.NoError(t, err) |
| 35 | + return rec != nil |
| 36 | +} |
| 37 | + |
| 38 | +func TestPruneActivitiesToSize_RemovesOldestUntilUnderBudget(t *testing.T) { |
| 39 | + m, cleanup := setupTestStorageForActivity(t) |
| 40 | + defer cleanup() |
| 41 | + |
| 42 | + // 10 records × ~10KB ≈ 100KB total. |
| 43 | + saveSizedActivities(t, m, 10, 10*1024) |
| 44 | + budget := int64(45 * 1024) // only the newest few fit |
| 45 | + |
| 46 | + deleted, err := m.PruneActivitiesToSize(budget) |
| 47 | + require.NoError(t, err) |
| 48 | + assert.Greater(t, deleted, 0) |
| 49 | + |
| 50 | + // Oldest pruned, newest retained. |
| 51 | + assert.False(t, exists(t, m, "id-00"), "oldest record should be pruned") |
| 52 | + assert.True(t, exists(t, m, "id-09"), "newest record must be retained") |
| 53 | + |
| 54 | + // Idempotent: a second pass deletes nothing more. |
| 55 | + again, err := m.PruneActivitiesToSize(budget) |
| 56 | + require.NoError(t, err) |
| 57 | + assert.Equal(t, 0, again, "second pass should be a no-op") |
| 58 | +} |
| 59 | + |
| 60 | +func TestPruneActivitiesToSize_AlreadyUnderBudget_NoOp(t *testing.T) { |
| 61 | + m, cleanup := setupTestStorageForActivity(t) |
| 62 | + defer cleanup() |
| 63 | + saveSizedActivities(t, m, 5, 1024) // ~5KB total |
| 64 | + |
| 65 | + deleted, err := m.PruneActivitiesToSize(10 * 1024 * 1024) // 10MB budget |
| 66 | + require.NoError(t, err) |
| 67 | + assert.Equal(t, 0, deleted) |
| 68 | + for i := 0; i < 5; i++ { |
| 69 | + assert.True(t, exists(t, m, fmt.Sprintf("id-%02d", i))) |
| 70 | + } |
| 71 | +} |
| 72 | + |
| 73 | +func TestPruneActivitiesToSize_KeepsNewestEvenIfOverBudget(t *testing.T) { |
| 74 | + m, cleanup := setupTestStorageForActivity(t) |
| 75 | + defer cleanup() |
| 76 | + // 5 records × 10KB; budget smaller than a single record. |
| 77 | + saveSizedActivities(t, m, 5, 10*1024) |
| 78 | + |
| 79 | + deleted, err := m.PruneActivitiesToSize(1024) // 1KB < one record |
| 80 | + require.NoError(t, err) |
| 81 | + assert.Equal(t, 4, deleted, "all but the newest should be deleted") |
| 82 | + |
| 83 | + // Only the newest survives — the log is never emptied. |
| 84 | + assert.True(t, exists(t, m, "id-04"), "newest must survive") |
| 85 | + for i := 0; i < 4; i++ { |
| 86 | + assert.Falsef(t, exists(t, m, fmt.Sprintf("id-%02d", i)), "id-%02d should be pruned", i) |
| 87 | + } |
| 88 | +} |
| 89 | + |
| 90 | +func TestPruneActivitiesToSize_DisabledWhenZeroOrNegative(t *testing.T) { |
| 91 | + m, cleanup := setupTestStorageForActivity(t) |
| 92 | + defer cleanup() |
| 93 | + saveSizedActivities(t, m, 4, 10*1024) |
| 94 | + |
| 95 | + for _, budget := range []int64{0, -1} { |
| 96 | + deleted, err := m.PruneActivitiesToSize(budget) |
| 97 | + require.NoError(t, err) |
| 98 | + assert.Equalf(t, 0, deleted, "budget %d disables size pruning", budget) |
| 99 | + } |
| 100 | + for i := 0; i < 4; i++ { |
| 101 | + assert.True(t, exists(t, m, fmt.Sprintf("id-%02d", i))) |
| 102 | + } |
| 103 | +} |
0 commit comments