Check operation context on returned errors

[onechat] / memory_test.go

diff --git a/memory_test.go b/memory_test.go
index 426654a63667c27508eabc2b7f3ad3c31fb4da78..3930435aca7eb400699d43de8cd0c7c3f49f1192 100644 (file)
--- a/memory_test.go
+++ b/memory_test.go
@@ -1,10 +1,16 @@
 package main
 
 import (
 package main
 
 import (

+       "context"

        "fmt"
        "fmt"

+       "log"
+       "maps"

        "os"
        "path/filepath"
        "os"
        "path/filepath"

+       "slices"
+       "strings"

        "testing"
        "testing"

+       "time"

 
        "github.com/stretchr/testify/require"
 )
 
        "github.com/stretchr/testify/require"
 )


@@ -12,7 +18,7 @@ import (
 func TestMemoryRestartLoadsTheSameView(t *testing.T) {
        root := t.TempDir()
        clock := newTestClock()
 func TestMemoryRestartLoadsTheSameView(t *testing.T) {
        root := t.TempDir()
        clock := newTestClock()

-       memory, err := openMemory(t.Context(), root, clock.Now)
+       memory, err := openMemory(t.Context(), root, clock.Now, newFakeSummarizer(), discardLog())

        require.NoError(t, err)
 
        memory.chat.high = 200
        require.NoError(t, err)
 
        memory.chat.high = 200


@@ -31,7 +37,7 @@ func TestMemoryRestartLoadsTheSameView(t *testing.T) {
        require.Less(t, len(memory.chat.lines), 100)
        require.NoError(t, memory.Close())
 
        require.Less(t, len(memory.chat.lines), 100)
        require.NoError(t, memory.Close())
 

-       memory, err = openMemory(t.Context(), root, clock.Now)
+       memory, err = openMemory(t.Context(), root, clock.Now, newFakeSummarizer(), discardLog())

        require.NoError(t, err)
 
        defer func() {
        require.NoError(t, err)
 
        defer func() {


@@ -46,8 +52,35 @@ func TestMemoryRestartLoadsTheSameView(t *testing.T) {
        require.Equal(t, saved, loaded)
 }
 
        require.Equal(t, saved, loaded)
 }
 

+func TestMemoryOpensAfterACrashBeforeTheViewSave(t *testing.T) {
+       root := t.TempDir()
+       clock := newTestClock()
+       memory, err := openMemory(t.Context(), root, clock.Now, newFakeSummarizer(), discardLog())
+       require.NoError(t, err)
+
+       _, err = memory.Append("user", "first")
+       require.NoError(t, err)
+
+       saved, err := os.ReadFile(filepath.Join(root, "view.json"))
+       require.NoError(t, err)
+
+       _, err = memory.Append("user", "second")
+       require.NoError(t, err)
+
+       view := memory.View()
+       require.NoError(t, memory.Close())
+
+       // The log and the tree hold message 1, but the view does not.
+       require.NoError(t, os.WriteFile(filepath.Join(root, "view.json"), saved, 0o600))
+
+       memory, err = openMemory(t.Context(), root, clock.Now, newFakeSummarizer(), discardLog())
+       require.NoError(t, err)
+       require.Equal(t, view, memory.View())
+       require.NoError(t, memory.Close())
+}
+

 func TestMemoryKeepsTheCompactionViewCoarser(t *testing.T) {
 func TestMemoryKeepsTheCompactionViewCoarser(t *testing.T) {

-       memory, err := openMemory(t.Context(), t.TempDir(), newTestClock().Now)
+       memory, err := openMemory(t.Context(), t.TempDir(), newTestClock().Now, newFakeSummarizer(), discardLog())

        require.NoError(t, err)
 
        defer func() {
        require.NoError(t, err)
 
        defer func() {


@@ -75,8 +108,200 @@ func TestMemoryKeepsTheCompactionViewCoarser(t *testing.T) {
        }
 }
 
        }
 }
 

+func TestMemoryRunsAtMostEightBuilds(t *testing.T) {
+       summarizer := newFakeSummarizer()
+       summarizer.gate = make(chan struct{})
+       memory, err := openMemory(t.Context(), t.TempDir(), newTestClock().Now, summarizer, discardLog())
+       require.NoError(t, err)
+
+       defer func() {
+               require.NoError(t, memory.Close())
+       }()
+
+       for n := range 12 {
+               _, err = memory.Append("echo", memoryLongText(n))
+               require.NoError(t, err)
+       }
+
+       memory.mutex.Lock()
+       running := slices.Collect(maps.Keys(memory.queue.running))
+       memory.mutex.Unlock()
+       slices.SortFunc(running, func(a, b NodeKey) int {
+               return a.index - b.index
+       })
+       require.Equal(t, []NodeKey{{0, 0}, {0, 1}, {0, 2}, {0, 3}, {0, 4}, {0, 5}, {0, 6}, {0, 7}}, running)
+
+       require.Eventually(t, func() bool {
+               return summarizer.Active() == buildMax
+       }, 10*time.Second, time.Millisecond)
+
+       close(summarizer.gate)
+       require.NoError(t, memory.Wait(t.Context()))
+       require.Equal(t, buildMax, summarizer.peak)
+
+       for level, count := range []int{12, 6, 3, 1} {
+               for index := range count {
+                       require.True(t, treeIsBuilt(t, memory.tree, level, index), "node (%d, %d)", level, index)
+               }
+       }
+}
+
+func TestMemorySettleWaitsForThePendingLines(t *testing.T) {
+       summarizer := newFakeSummarizer()
+       summarizer.gate = make(chan struct{})
+       memory, err := openMemory(t.Context(), t.TempDir(), newTestClock().Now, summarizer, discardLog())
+       require.NoError(t, err)
+
+       defer func() {
+               require.NoError(t, memory.Close())
+       }()
+
+       _, err = memory.Append("echo", memoryLongText(0))
+       require.NoError(t, err)
+       require.Contains(t, memory.View(), viewPending)
+
+       settled := make(chan error, 1)
+       go func() {
+               settled <- memory.Settle(t.Context())
+       }()
+
+       select {
+       case err = <-settled:
+               require.Fail(t, "Settle returned before the build", "error: %v", err)
+       case <-time.After(50 * time.Millisecond):
+       }
+
+       close(summarizer.gate)
+       require.NoError(t, <-settled)
+       require.NotContains(t, memory.View(), viewPending)
+}
+
+func TestMemorySettleFailsWhileABuildFailed(t *testing.T) {
+       summarizer := newFakeSummarizer()
+       summarizer.SetFailing(true)
+       logs := &SyncBuffer{}
+       memory, err := openMemory(t.Context(), t.TempDir(), newTestClock().Now, summarizer, log.New(logs, "", 0))
+       require.NoError(t, err)
+
+       defer func() {
+               require.NoError(t, memory.Close())
+       }()
+
+       _, err = memory.Append("echo", memoryLongText(0))
+       require.NoError(t, err)
+       require.ErrorIs(t, memory.Settle(t.Context()), errMemoryUnsettled)
+       require.Equal(t, "summary failed, retry at next message: node (0, 0): fake failure\n", logs.String())
+
+       summarizer.SetFailing(false)
+       _, err = memory.Append("user", "retry")
+       require.NoError(t, err)
+       require.NoError(t, memory.Settle(t.Context()))
+}
+
+func TestMemorySettleEndsWithItsContext(t *testing.T) {
+       summarizer := newFakeSummarizer()
+       summarizer.gate = make(chan struct{})
+       memory, err := openMemory(t.Context(), t.TempDir(), newTestClock().Now, summarizer, discardLog())
+       require.NoError(t, err)
+
+       defer func() {
+               require.NoError(t, memory.Close())
+       }()
+
+       _, err = memory.Append("echo", memoryLongText(0))
+       require.NoError(t, err)
+
+       ctx, cancel := context.WithTimeout(t.Context(), 50*time.Millisecond)
+       defer cancel()
+
+       require.ErrorIs(t, memory.Settle(ctx), context.DeadlineExceeded)
+}
+
+func TestMemoryStartsALeafAfterFewerThanEightUnbuilt(t *testing.T) {
+       summarizer := newFakeSummarizer()
+       summarizer.SetFailing(true)
+       memory, err := openMemory(t.Context(), t.TempDir(), newTestClock().Now, summarizer, discardLog())
+       require.NoError(t, err)
+
+       defer func() {
+               require.NoError(t, memory.Close())
+       }()
+
+       for n := range 9 {
+               _, err = memory.Append("echo", memoryLongText(n))
+               require.NoError(t, err)
+       }
+
+       require.NoError(t, memory.Wait(t.Context()))
+       for _, task := range summarizer.Tasks() {
+               require.Less(t, task.Index, 8)
+       }
+
+       require.False(t, treeIsBuilt(t, memory.tree, 0, 8))
+       summarizer.SetFailing(false)
+
+       _, err = memory.Append("user", "retry")
+       require.NoError(t, err)
+       require.NoError(t, memory.Wait(t.Context()))
+       require.True(t, treeIsBuilt(t, memory.tree, 0, 8))
+       require.True(t, treeIsBuilt(t, memory.tree, 3, 0))
+}
+
+func TestMemoryMergesAfterBothHalves(t *testing.T) {
+       summarizer := newFakeSummarizer()
+       memory, err := openMemory(t.Context(), t.TempDir(), newTestClock().Now, summarizer, discardLog())
+       require.NoError(t, err)
+
+       defer func() {
+               require.NoError(t, memory.Close())
+       }()
+
+       _, err = memory.Append("echo", memoryLongText(0))
+       require.NoError(t, err)
+       require.NoError(t, memory.Wait(t.Context()))
+
+       _, err = memory.Append("echo", memoryLongText(1))
+       require.NoError(t, err)
+       require.NoError(t, memory.Wait(t.Context()))
+       require.Len(t, summarizer.Tasks(), 3)
+
+       left := treeReadNode(t, memory.tree, 0, 0)
+       right := treeReadNode(t, memory.tree, 0, 1)
+       merge := summarizer.Tasks()[2]
+       require.Equal(t, []string{left.Text, right.Text}, merge.Input)
+       require.Equal(t, "<chat>\n0+1|"+left.Text+"\n1+1|"+right.Text+"\n</chat>", merge.Context)
+       require.True(t, treeIsBuilt(t, memory.tree, 1, 0))
+}
+
+func TestMemoryRestartRestoresTheQueue(t *testing.T) {
+       root := t.TempDir()
+       clock := newTestClock()
+       summarizer := newFakeSummarizer()
+       summarizer.SetFailing(true)
+       memory, err := openMemory(t.Context(), root, clock.Now, summarizer, discardLog())
+       require.NoError(t, err)
+
+       for n := range 4 {
+               _, err = memory.Append("echo", memoryLongText(n))
+               require.NoError(t, err)
+       }
+
+       require.NoError(t, memory.Wait(t.Context()))
+       require.NoError(t, memory.Close())
+
+       memory, err = openMemory(t.Context(), root, clock.Now, newFakeSummarizer(), discardLog())
+       require.NoError(t, err)
+
+       defer func() {
+               require.NoError(t, memory.Close())
+       }()
+
+       require.NoError(t, memory.Wait(t.Context()))
+       require.True(t, treeIsBuilt(t, memory.tree, 2, 0))
+}
+

 func TestMemoryBuildsFreeParents(t *testing.T) {
 func TestMemoryBuildsFreeParents(t *testing.T) {

-       memory, err := openMemory(t.Context(), t.TempDir(), newTestClock().Now)
+       memory, err := openMemory(t.Context(), t.TempDir(), newTestClock().Now, newFakeSummarizer(), discardLog())

        require.NoError(t, err)
 
        defer func() {
        require.NoError(t, err)
 
        defer func() {


@@ -88,29 +313,40 @@ func TestMemoryBuildsFreeParents(t *testing.T) {
                require.NoError(t, err)
        }
 
                require.NoError(t, err)
        }
 

-       node, built := memory.tree.Node(2, 0)
-       require.True(t, built)
+       node := treeReadNode(t, memory.tree, 2, 0)

        require.Equal(t, "user: a\nuser: b\nuser: c\nuser: d", node.Text)
 }
 
 func TestMemoryKeepsALongMessageOutOfTheView(t *testing.T) {
        require.Equal(t, "user: a\nuser: b\nuser: c\nuser: d", node.Text)
 }
 
 func TestMemoryKeepsALongMessageOutOfTheView(t *testing.T) {

-       memory, err := openMemory(t.Context(), t.TempDir(), newTestClock().Now)
+       summarizer := newFakeSummarizer()
+       summarizer.SetFailing(true)
+       memory, err := openMemory(t.Context(), t.TempDir(), newTestClock().Now, summarizer, discardLog())

        require.NoError(t, err)
 
        defer func() {
                require.NoError(t, memory.Close())
        }()
 
        require.NoError(t, err)
 
        defer func() {
                require.NoError(t, memory.Close())
        }()
 

-       _, err = memory.Append("echo", string(make([]byte, nodeBytes)))
+       _, err = memory.Append("echo", memoryLongText(0))
+       require.NoError(t, err)
+       require.NoError(t, memory.Wait(t.Context()))
+       require.Equal(t, "<chat>\n0+1|"+viewPending+"\n</chat>", memory.View())
+
+       summarizer.SetFailing(false)
+
+       _, err = memory.Append("user", "again")

        require.NoError(t, err)
        require.NoError(t, err)

+       require.NoError(t, memory.Wait(t.Context()))

 
 

-       require.Equal(t, "<chat>\n0+1|"+pending+"\n</chat>", memory.View())
+       task, err := memory.task(NodeKey{0, 0})
+       require.NoError(t, err)
+       require.Equal(t, "<chat>\n0+1|"+fakeLine(task)+"\n1+1|user: again\n</chat>", memory.View())

 }
 
 func TestMemoryCoversMessagesAfterTheSavedView(t *testing.T) {
        root := t.TempDir()
        clock := newTestClock()
 }
 
 func TestMemoryCoversMessagesAfterTheSavedView(t *testing.T) {
        root := t.TempDir()
        clock := newTestClock()

-       memory, err := openMemory(t.Context(), root, clock.Now)
+       memory, err := openMemory(t.Context(), root, clock.Now, newFakeSummarizer(), discardLog())

        require.NoError(t, err)
 
        _, err = memory.Append("user", "a")
        require.NoError(t, err)
 
        _, err = memory.Append("user", "a")


@@ -120,7 +356,7 @@ func TestMemoryCoversMessagesAfterTheSavedView(t *testing.T) {
        require.NoError(t, err)
        require.NoError(t, memory.Close())
 
        require.NoError(t, err)
        require.NoError(t, memory.Close())
 

-       memory, err = openMemory(t.Context(), root, clock.Now)
+       memory, err = openMemory(t.Context(), root, clock.Now, newFakeSummarizer(), discardLog())

        require.NoError(t, err)
 
        defer func() {
        require.NoError(t, err)
 
        defer func() {


@@ -133,7 +369,7 @@ func TestMemoryCoversMessagesAfterTheSavedView(t *testing.T) {
 func TestMemoryRefusesToRebuildAMissingView(t *testing.T) {
        root := t.TempDir()
        clock := newTestClock()
 func TestMemoryRefusesToRebuildAMissingView(t *testing.T) {
        root := t.TempDir()
        clock := newTestClock()

-       memory, err := openMemory(t.Context(), root, clock.Now)
+       memory, err := openMemory(t.Context(), root, clock.Now, newFakeSummarizer(), discardLog())

        require.NoError(t, err)
 
        _, err = memory.Append("user", "a")
        require.NoError(t, err)
 
        _, err = memory.Append("user", "a")


@@ -141,6 +377,11 @@ func TestMemoryRefusesToRebuildAMissingView(t *testing.T) {
        require.NoError(t, memory.Close())
        require.NoError(t, os.Remove(filepath.Join(root, "view.json")))
 
        require.NoError(t, memory.Close())
        require.NoError(t, os.Remove(filepath.Join(root, "view.json")))
 

-       _, err = openMemory(t.Context(), root, clock.Now)
-       require.ErrorIs(t, err, errNoView)
+       _, err = openMemory(t.Context(), root, clock.Now, newFakeSummarizer(), discardLog())
+       require.ErrorIs(t, err, errMemoryNoView)
+}
+
+// memoryLongText returns a text that is too long for its own node.
+func memoryLongText(n int) string {
+       return fmt.Sprintf("long %d ", n) + strings.Repeat("x", nodeBytes)

 }
 }