X-Git-Url: https://git.giacomolli.uk/onechat/blobdiff_plain/d4dc9a3b10a6092d18fb30cfb36526556aa33227..cc5d2d2d302e810c6d9b862035e9f8de91941390:/memory_test.go

diff --git a/memory_test.go b/memory_test.go
index ddbb014..8c34149 100644
--- a/memory_test.go
+++ b/memory_test.go
@@ -116,7 +116,7 @@ func TestMemoryRunsAtMostEightBuilds(t *testing.T) {
 
 	for level, count := range []int{12, 6, 3, 1} {
 		for index := range count {
-			require.True(t, memory.tree.Built(level, index), "node (%d, %d)", level, index)
+			require.True(t, isBuilt(t, memory.tree, level, index), "node (%d, %d)", level, index)
 		}
 	}
 }
@@ -141,14 +141,14 @@ func TestMemoryStartsALeafAfterFewerThanEightUnbuilt(t *testing.T) {
 		require.Less(t, task.Index, 8)
 	}
 
-	require.False(t, memory.tree.Built(0, 8))
+	require.False(t, isBuilt(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, memory.tree.Built(0, 8))
-	require.True(t, memory.tree.Built(3, 0))
+	require.True(t, isBuilt(t, memory.tree, 0, 8))
+	require.True(t, isBuilt(t, memory.tree, 3, 0))
 }
 
 func TestMemoryMergesAfterBothHalves(t *testing.T) {
@@ -169,12 +169,12 @@ func TestMemoryMergesAfterBothHalves(t *testing.T) {
 	require.NoError(t, memory.Wait(t.Context()))
 	require.Len(t, summarizer.Tasks(), 3)
 
-	left, _ := memory.tree.Node(0, 0)
-	right, _ := memory.tree.Node(0, 1)
+	left := readNode(t, memory.tree, 0, 0)
+	right := readNode(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, memory.tree.Built(1, 0))
+	require.True(t, isBuilt(t, memory.tree, 1, 0))
 }
 
 func TestMemoryRestartRestoresTheQueue(t *testing.T) {
@@ -201,7 +201,7 @@ func TestMemoryRestartRestoresTheQueue(t *testing.T) {
 	}()
 
 	require.NoError(t, memory.Wait(t.Context()))
-	require.True(t, memory.tree.Built(2, 0))
+	require.True(t, isBuilt(t, memory.tree, 2, 0))
 }
 
 func TestMemoryBuildsFreeParents(t *testing.T) {
@@ -217,8 +217,7 @@ func TestMemoryBuildsFreeParents(t *testing.T) {
 		require.NoError(t, err)
 	}
 
-	node, built := memory.tree.Node(2, 0)
-	require.True(t, built)
+	node := readNode(t, memory.tree, 2, 0)
 	require.Equal(t, "user: a\nuser: b\nuser: c\nuser: d", node.Text)
 }
 
@@ -242,7 +241,10 @@ func TestMemoryKeepsALongMessageOutOfTheView(t *testing.T) {
 	_, err = memory.Append("user", "again")
 	require.NoError(t, err)
 	require.NoError(t, memory.Wait(t.Context()))
-	require.Equal(t, "<chat>\n0+1|"+fakeLine(memory.task(nodeKey{0, 0}))+"\n1+1|user: again\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) {