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

   1 package main
   2
   3 import (
   4         "os"
   5         "path/filepath"
   6         "strings"
   7         "testing"
   8
   9         "github.com/stretchr/testify/require"
  10 )
  11
  12 func TestFreeLeafHoldsAShortMessageWordForWord(t *testing.T) {
  13         message := Message{
  14                 ID:   7,
  15                 Kind: "user",
  16                 Text: strings.Repeat("a", nodeBytes-len("user: ")),
  17         }
  18
  19         node, free := freeLeaf(message)
  20         require.True(t, free)
  21         require.Equal(t, Node{
  22                 Level: 0,
  23                 Index: 7,
  24                 Text:  "user: " + message.Text,
  25                 Size:  nodeBytes,
  26         }, node)
  27
  28         message.Text += "a"
  29         _, free = freeLeaf(message)
  30         require.False(t, free)
  31 }
  32
  33 func TestFreeMergeJoinsShortSiblings(t *testing.T) {
  34         left := testNode(1, 2, "left")
  35         right := testNode(1, 3, "right")
  36
  37         node, free := freeMerge(left, right)
  38         require.True(t, free)
  39         require.Equal(t, testNode(2, 1, "left\nright"), node)
  40
  41         right = testNode(1, 3, strings.Repeat("r", nodeBytes-len("left")))
  42         _, free = freeMerge(left, right)
  43         require.False(t, free)
  44 }
  45
  46 func TestTreeBuildsEachNodeOnce(t *testing.T) {
  47         tree, err := openTree(t.Context(), t.TempDir(), 2, newTestClock().Now)
  48         require.NoError(t, err)
  49
  50         defer func() {
  51                 require.NoError(t, tree.Close())
  52         }()
  53
  54         require.NoError(t, tree.Add(testNode(0, 0, "user: a")))
  55         require.ErrorContains(t, tree.Add(testNode(0, 0, "user: b")), "already built")
  56         require.ErrorContains(t, tree.Add(testNode(1, 0, "a b")), "child that is not built")
  57
  58         require.NoError(t, tree.Add(testNode(0, 1, "reply: b")))
  59         require.NoError(t, tree.Add(testNode(1, 0, "a b")))
  60
  61         require.Equal(t, "user: a", readNode(t, tree, 0, 0).Text)
  62 }
  63
  64 func TestTreeReopenLoadsTheSameNodes(t *testing.T) {
  65         root := t.TempDir()
  66         clock := newTestClock()
  67         tree, err := openTree(t.Context(), root, 2, clock.Now)
  68         require.NoError(t, err)
  69
  70         nodes := []Node{
  71                 testNode(0, 1, "reply: b"),
  72                 testNode(0, 0, "user: a"),
  73                 testNode(1, 0, "user: a\nreply: b"),
  74         }
  75         for _, node := range nodes {
  76                 require.NoError(t, tree.Add(node))
  77         }
  78
  79         require.NoError(t, tree.Close())
  80
  81         data, err := os.ReadFile(filepath.Join(root, "tree", "2026-01-01.jsonl"))
  82         require.NoError(t, err)
  83         require.Equal(t, `{"l":0,"i":1,"text":"reply: b","size":8}`+"\n", strings.SplitAfter(string(data), "\n")[0])
  84
  85         tree, err = openTree(t.Context(), root, 2, clock.Now)
  86         require.NoError(t, err)
  87
  88         defer func() {
  89                 require.NoError(t, tree.Close())
  90         }()
  91
  92         require.Equal(t, len(nodes), tree.Count())
  93         for _, node := range nodes {
  94                 require.Equal(t, node, readNode(t, tree, node.Level, node.Index))
  95         }
  96 }
  97
  98 func TestTreeBuildsAMissingIndexAgain(t *testing.T) {
  99         root := t.TempDir()
 100         clock := newTestClock()
 101         tree, err := openTree(t.Context(), root, 2, clock.Now)
 102         require.NoError(t, err)
 103
 104         nodes := []Node{
 105                 testNode(0, 1, "reply: b"),
 106                 testNode(0, 0, "user: a"),
 107                 testNode(1, 0, "user: a\nreply: b"),
 108         }
 109         for _, node := range nodes {
 110                 require.NoError(t, tree.Add(node))
 111         }
 112
 113         require.NoError(t, tree.Close())
 114         require.NoError(t, os.RemoveAll(filepath.Join(root, "index")))
 115
 116         _, err = openTree(t.Context(), root, 1, clock.Now)
 117         require.ErrorContains(t, err, "covers messages after the log")
 118         require.NoError(t, os.RemoveAll(filepath.Join(root, "index")))
 119
 120         tree, err = openTree(t.Context(), root, 2, clock.Now)
 121         require.NoError(t, err)
 122
 123         defer func() {
 124                 require.NoError(t, tree.Close())
 125         }()
 126
 127         require.Equal(t, len(nodes), tree.Count())
 128         for _, node := range nodes {
 129                 require.Equal(t, node, readNode(t, tree, node.Level, node.Index))
 130         }
 131 }
 132
 133 func TestTreeIndexesTheNodesAfterACrash(t *testing.T) {
 134         root := t.TempDir()
 135         clock := newTestClock()
 136         tree, err := openTree(t.Context(), root, 2, clock.Now)
 137         require.NoError(t, err)
 138
 139         require.NoError(t, tree.Add(testNode(0, 0, "user: a")))
 140         require.NoError(t, tree.Add(testNode(0, 1, "reply: b")))
 141         require.NoError(t, tree.Close())
 142
 143         // A crash after the line write and in the middle of the index writes.
 144         order := filepath.Join(root, "index", "tree.idx")
 145         require.NoError(t, os.Truncate(order, treeRecordBytes+5))
 146         require.NoError(t, os.Truncate(filepath.Join(root, "index", "tree-0.idx"), positionBytes))
 147
 148         tree, err = openTree(t.Context(), root, 2, clock.Now)
 149         require.NoError(t, err)
 150
 151         defer func() {
 152                 require.NoError(t, tree.Close())
 153         }()
 154
 155         require.Equal(t, 2, tree.Count())
 156         require.Equal(t, "reply: b", readNode(t, tree, 0, 1).Text)
 157         require.NoError(t, tree.Add(testNode(1, 0, "user: a\nreply: b")))
 158 }
 159
 160 // isBuilt tells if the node (l, i) of tree is built.
 161 func isBuilt(t *testing.T, tree *Tree, level, index int) bool {
 162         t.Helper()
 163
 164         built, err := tree.Built(level, index)
 165         require.NoError(t, err)
 166
 167         return built
 168 }
 169
 170 // readNode reads the node (l, i) of tree, which must be built.
 171 func readNode(t *testing.T, tree *Tree, level, index int) Node {
 172         t.Helper()
 173
 174         node, built, err := tree.Node(level, index)
 175         require.NoError(t, err)
 176         require.True(t, built, "node (%d, %d)", level, index)
 177
 178         return node
 179 }
 180
 181 func testNode(level, index int, text string) Node {
 182         return Node{
 183                 Level: level,
 184                 Index: index,
 185                 Text:  text,
 186                 Size:  len(text),
 187         }
 188 }