git.giacomolli.uk - onechat/blob - memory_test.go

   1 package main
   2
   3 import (
   4         "fmt"
   5         "maps"
   6         "os"
   7         "path/filepath"
   8         "slices"
   9         "strings"
  10         "testing"
  11         "time"
  12
  13         "github.com/stretchr/testify/require"
  14 )
  15
  16 func TestMemoryRestartLoadsTheSameView(t *testing.T) {
  17         root := t.TempDir()
  18         clock := newTestClock()
  19         memory, err := openMemory(t.Context(), root, clock.Now, newFakeSummarizer())
  20         require.NoError(t, err)
  21
  22         memory.chat.high = 200
  23         memory.chat.low = 100
  24         memory.cview.high = 80
  25         memory.cview.low = 40
  26         for n := range 100 {
  27                 _, err = memory.Append("user", fmt.Sprint(n))
  28                 require.NoError(t, err)
  29         }
  30
  31         view := memory.View()
  32         context := memory.Context(0, 100)
  33         saved, err := os.ReadFile(filepath.Join(root, "view.json"))
  34         require.NoError(t, err)
  35         require.Less(t, len(memory.chat.lines), 100)
  36         require.NoError(t, memory.Close())
  37
  38         memory, err = openMemory(t.Context(), root, clock.Now, newFakeSummarizer())
  39         require.NoError(t, err)
  40
  41         defer func() {
  42                 require.NoError(t, memory.Close())
  43         }()
  44
  45         require.Equal(t, view, memory.View())
  46         require.Equal(t, context, memory.Context(0, 100))
  47
  48         loaded, err := os.ReadFile(filepath.Join(root, "view.json"))
  49         require.NoError(t, err)
  50         require.Equal(t, saved, loaded)
  51 }
  52
  53 func TestMemoryOpensAfterACrashBeforeTheViewSave(t *testing.T) {
  54         root := t.TempDir()
  55         clock := newTestClock()
  56         memory, err := openMemory(t.Context(), root, clock.Now, newFakeSummarizer())
  57         require.NoError(t, err)
  58
  59         _, err = memory.Append("user", "first")
  60         require.NoError(t, err)
  61
  62         saved, err := os.ReadFile(filepath.Join(root, "view.json"))
  63         require.NoError(t, err)
  64
  65         _, err = memory.Append("user", "second")
  66         require.NoError(t, err)
  67
  68         view := memory.View()
  69         require.NoError(t, memory.Close())
  70
  71         // The log and the tree hold message 1, but the view does not.
  72         require.NoError(t, os.WriteFile(filepath.Join(root, "view.json"), saved, 0o600))
  73
  74         memory, err = openMemory(t.Context(), root, clock.Now, newFakeSummarizer())
  75         require.NoError(t, err)
  76         require.Equal(t, view, memory.View())
  77         require.NoError(t, memory.Close())
  78 }
  79
  80 func TestMemoryKeepsTheCompactionViewCoarser(t *testing.T) {
  81         memory, err := openMemory(t.Context(), t.TempDir(), newTestClock().Now, newFakeSummarizer())
  82         require.NoError(t, err)
  83
  84         defer func() {
  85                 require.NoError(t, memory.Close())
  86         }()
  87
  88         memory.chat.high = 200
  89         memory.chat.low = 100
  90         memory.cview.high = 80
  91         memory.cview.low = 40
  92         for n := range 64 {
  93                 _, err = memory.Append("user", fmt.Sprint(n%10))
  94                 require.NoError(t, err)
  95
  96                 coarse := 0
  97                 for _, line := range memory.chat.lines {
  98                         outer := memory.cview.lines[coarse]
  99                         for nodeEnd(outer.level, outer.index) <= line.index<<line.level {
 100                                 coarse++
 101                                 outer = memory.cview.lines[coarse]
 102                         }
 103
 104                         require.GreaterOrEqual(t, outer.level, line.level, "message %d", n)
 105                 }
 106         }
 107 }
 108
 109 func TestMemoryRunsAtMostEightBuilds(t *testing.T) {
 110         summarizer := newFakeSummarizer()
 111         summarizer.gate = make(chan struct{})
 112         memory, err := openMemory(t.Context(), t.TempDir(), newTestClock().Now, summarizer)
 113         require.NoError(t, err)
 114
 115         defer func() {
 116                 require.NoError(t, memory.Close())
 117         }()
 118
 119         for n := range 12 {
 120                 _, err = memory.Append("echo", memoryLongText(n))
 121                 require.NoError(t, err)
 122         }
 123
 124         memory.mutex.Lock()
 125         running := slices.Collect(maps.Keys(memory.queue.running))
 126         memory.mutex.Unlock()
 127         slices.SortFunc(running, func(a, b NodeKey) int {
 128                 return a.index - b.index
 129         })
 130         require.Equal(t, []NodeKey{{0, 0}, {0, 1}, {0, 2}, {0, 3}, {0, 4}, {0, 5}, {0, 6}, {0, 7}}, running)
 131
 132         require.Eventually(t, func() bool {
 133                 return summarizer.Active() == buildMax
 134         }, 10*time.Second, time.Millisecond)
 135
 136         close(summarizer.gate)
 137         require.NoError(t, memory.Wait(t.Context()))
 138         require.Equal(t, buildMax, summarizer.peak)
 139
 140         for level, count := range []int{12, 6, 3, 1} {
 141                 for index := range count {
 142                         require.True(t, treeIsBuilt(t, memory.tree, level, index), "node (%d, %d)", level, index)
 143                 }
 144         }
 145 }
 146
 147 func TestMemoryStartsALeafAfterFewerThanEightUnbuilt(t *testing.T) {
 148         summarizer := newFakeSummarizer()
 149         summarizer.SetFailing(true)
 150         memory, err := openMemory(t.Context(), t.TempDir(), newTestClock().Now, summarizer)
 151         require.NoError(t, err)
 152
 153         defer func() {
 154                 require.NoError(t, memory.Close())
 155         }()
 156
 157         for n := range 9 {
 158                 _, err = memory.Append("echo", memoryLongText(n))
 159                 require.NoError(t, err)
 160         }
 161
 162         require.NoError(t, memory.Wait(t.Context()))
 163         for _, task := range summarizer.Tasks() {
 164                 require.Less(t, task.Index, 8)
 165         }
 166
 167         require.False(t, treeIsBuilt(t, memory.tree, 0, 8))
 168         summarizer.SetFailing(false)
 169
 170         _, err = memory.Append("user", "retry")
 171         require.NoError(t, err)
 172         require.NoError(t, memory.Wait(t.Context()))
 173         require.True(t, treeIsBuilt(t, memory.tree, 0, 8))
 174         require.True(t, treeIsBuilt(t, memory.tree, 3, 0))
 175 }
 176
 177 func TestMemoryMergesAfterBothHalves(t *testing.T) {
 178         summarizer := newFakeSummarizer()
 179         memory, err := openMemory(t.Context(), t.TempDir(), newTestClock().Now, summarizer)
 180         require.NoError(t, err)
 181
 182         defer func() {
 183                 require.NoError(t, memory.Close())
 184         }()
 185
 186         _, err = memory.Append("echo", memoryLongText(0))
 187         require.NoError(t, err)
 188         require.NoError(t, memory.Wait(t.Context()))
 189
 190         _, err = memory.Append("echo", memoryLongText(1))
 191         require.NoError(t, err)
 192         require.NoError(t, memory.Wait(t.Context()))
 193         require.Len(t, summarizer.Tasks(), 3)
 194
 195         left := treeReadNode(t, memory.tree, 0, 0)
 196         right := treeReadNode(t, memory.tree, 0, 1)
 197         merge := summarizer.Tasks()[2]
 198         require.Equal(t, []string{left.Text, right.Text}, merge.Input)
 199         require.Equal(t, "<chat>\n0+1|"+left.Text+"\n1+1|"+right.Text+"\n</chat>", merge.Context)
 200         require.True(t, treeIsBuilt(t, memory.tree, 1, 0))
 201 }
 202
 203 func TestMemoryRestartRestoresTheQueue(t *testing.T) {
 204         root := t.TempDir()
 205         clock := newTestClock()
 206         summarizer := newFakeSummarizer()
 207         summarizer.SetFailing(true)
 208         memory, err := openMemory(t.Context(), root, clock.Now, summarizer)
 209         require.NoError(t, err)
 210
 211         for n := range 4 {
 212                 _, err = memory.Append("echo", memoryLongText(n))
 213                 require.NoError(t, err)
 214         }
 215
 216         require.NoError(t, memory.Wait(t.Context()))
 217         require.NoError(t, memory.Close())
 218
 219         memory, err = openMemory(t.Context(), root, clock.Now, newFakeSummarizer())
 220         require.NoError(t, err)
 221
 222         defer func() {
 223                 require.NoError(t, memory.Close())
 224         }()
 225
 226         require.NoError(t, memory.Wait(t.Context()))
 227         require.True(t, treeIsBuilt(t, memory.tree, 2, 0))
 228 }
 229
 230 func TestMemoryBuildsFreeParents(t *testing.T) {
 231         memory, err := openMemory(t.Context(), t.TempDir(), newTestClock().Now, newFakeSummarizer())
 232         require.NoError(t, err)
 233
 234         defer func() {
 235                 require.NoError(t, memory.Close())
 236         }()
 237
 238         for _, text := range []string{"a", "b", "c", "d"} {
 239                 _, err = memory.Append("user", text)
 240                 require.NoError(t, err)
 241         }
 242
 243         node := treeReadNode(t, memory.tree, 2, 0)
 244         require.Equal(t, "user: a\nuser: b\nuser: c\nuser: d", node.Text)
 245 }
 246
 247 func TestMemoryKeepsALongMessageOutOfTheView(t *testing.T) {
 248         summarizer := newFakeSummarizer()
 249         summarizer.SetFailing(true)
 250         memory, err := openMemory(t.Context(), t.TempDir(), newTestClock().Now, summarizer)
 251         require.NoError(t, err)
 252
 253         defer func() {
 254                 require.NoError(t, memory.Close())
 255         }()
 256
 257         _, err = memory.Append("echo", memoryLongText(0))
 258         require.NoError(t, err)
 259         require.NoError(t, memory.Wait(t.Context()))
 260         require.Equal(t, "<chat>\n0+1|"+viewPending+"\n</chat>", memory.View())
 261
 262         summarizer.SetFailing(false)
 263
 264         _, err = memory.Append("user", "again")
 265         require.NoError(t, err)
 266         require.NoError(t, memory.Wait(t.Context()))
 267
 268         task, err := memory.task(NodeKey{0, 0})
 269         require.NoError(t, err)
 270         require.Equal(t, "<chat>\n0+1|"+fakeLine(task)+"\n1+1|user: again\n</chat>", memory.View())
 271 }
 272
 273 func TestMemoryCoversMessagesAfterTheSavedView(t *testing.T) {
 274         root := t.TempDir()
 275         clock := newTestClock()
 276         memory, err := openMemory(t.Context(), root, clock.Now, newFakeSummarizer())
 277         require.NoError(t, err)
 278
 279         _, err = memory.Append("user", "a")
 280         require.NoError(t, err)
 281
 282         _, err = memory.log.Append("reply", "b")
 283         require.NoError(t, err)
 284         require.NoError(t, memory.Close())
 285
 286         memory, err = openMemory(t.Context(), root, clock.Now, newFakeSummarizer())
 287         require.NoError(t, err)
 288
 289         defer func() {
 290                 require.NoError(t, memory.Close())
 291         }()
 292
 293         require.Equal(t, "<chat>\n0+1|user: a\n1+1|reply: b\n</chat>", memory.View())
 294 }
 295
 296 func TestMemoryRefusesToRebuildAMissingView(t *testing.T) {
 297         root := t.TempDir()
 298         clock := newTestClock()
 299         memory, err := openMemory(t.Context(), root, clock.Now, newFakeSummarizer())
 300         require.NoError(t, err)
 301
 302         _, err = memory.Append("user", "a")
 303         require.NoError(t, err)
 304         require.NoError(t, memory.Close())
 305         require.NoError(t, os.Remove(filepath.Join(root, "view.json")))
 306
 307         _, err = openMemory(t.Context(), root, clock.Now, newFakeSummarizer())
 308         require.ErrorIs(t, err, errMemoryNoView)
 309 }
 310
 311 // memoryLongText returns a text that is too long for its own node.
 312 func memoryLongText(n int) string {
 313         return fmt.Sprintf("long %d ", n) + strings.Repeat("x", nodeBytes)
 314 }