restic/internal/walker/rewriter_test.go
Alex Johnson 3bf2927006 Update snapshot summary on rewrite
Signed-off-by: Alex Johnson <hello@alex-johnson.net>
2024-07-16 12:06:50 -05:00

420 lines
9.5 KiB
Go

package walker
import (
"context"
"fmt"
"testing"
"github.com/pkg/errors"
"github.com/restic/restic/internal/restic"
"github.com/restic/restic/internal/test"
)
// WritableTreeMap also support saving
type WritableTreeMap struct {
TreeMap
}
func (t WritableTreeMap) SaveBlob(_ context.Context, tpe restic.BlobType, buf []byte, id restic.ID, _ bool) (newID restic.ID, known bool, size int, err error) {
if tpe != restic.TreeBlob {
return restic.ID{}, false, 0, errors.New("can only save trees")
}
if id.IsNull() {
id = restic.Hash(buf)
}
_, ok := t.TreeMap[id]
if ok {
return id, false, 0, nil
}
t.TreeMap[id] = append([]byte{}, buf...)
return id, true, len(buf), nil
}
func (t WritableTreeMap) Dump() {
for k, v := range t.TreeMap {
fmt.Printf("%v: %v", k, string(v))
}
}
type checkRewriteFunc func(t testing.TB) (rewriter *TreeRewriter, final func(testing.TB))
// checkRewriteItemOrder ensures that the order of the 'path' arguments is the one passed in as 'want'.
func checkRewriteItemOrder(want []string) checkRewriteFunc {
pos := 0
return func(t testing.TB) (rewriter *TreeRewriter, final func(testing.TB)) {
rewriter = NewTreeRewriter(RewriteOpts{
RewriteNode: func(node *restic.Node, path string) *restic.Node {
if pos >= len(want) {
t.Errorf("additional unexpected path found: %v", path)
return nil
}
if path != want[pos] {
t.Errorf("wrong path found, want %q, got %q", want[pos], path)
}
pos++
return node
},
})
final = func(t testing.TB) {
if pos != len(want) {
t.Errorf("not enough items returned, want %d, got %d", len(want), pos)
}
}
return rewriter, final
}
}
// checkRewriteSkips excludes nodes if path is in skipFor, it checks that rewriting proceeds in the correct order.
func checkRewriteSkips(skipFor map[string]struct{}, want []string, disableCache bool) checkRewriteFunc {
var pos int
return func(t testing.TB) (rewriter *TreeRewriter, final func(testing.TB)) {
rewriter = NewTreeRewriter(RewriteOpts{
RewriteNode: func(node *restic.Node, path string) *restic.Node {
if pos >= len(want) {
t.Errorf("additional unexpected path found: %v", path)
return nil
}
if path != want[pos] {
t.Errorf("wrong path found, want %q, got %q", want[pos], path)
}
pos++
_, skip := skipFor[path]
if skip {
return nil
}
return node
},
DisableNodeCache: disableCache,
})
final = func(t testing.TB) {
if pos != len(want) {
t.Errorf("not enough items returned, want %d, got %d", len(want), pos)
}
}
return rewriter, final
}
}
// checkIncreaseNodeSize modifies each node by changing its size.
func checkIncreaseNodeSize(increase uint64) checkRewriteFunc {
return func(t testing.TB) (rewriter *TreeRewriter, final func(testing.TB)) {
rewriter = NewTreeRewriter(RewriteOpts{
RewriteNode: func(node *restic.Node, path string) *restic.Node {
if node.Type == "file" {
node.Size += increase
}
return node
},
})
final = func(t testing.TB) {}
return rewriter, final
}
}
func TestRewriter(t *testing.T) {
var tests = []struct {
tree TestTree
newTree TestTree
check checkRewriteFunc
}{
{ // don't change
tree: TestTree{
"foo": TestFile{},
"subdir": TestTree{
"subfile": TestFile{},
},
},
check: checkRewriteItemOrder([]string{
"/foo",
"/subdir",
"/subdir/subfile",
}),
},
{ // exclude file
tree: TestTree{
"foo": TestFile{},
"subdir": TestTree{
"subfile": TestFile{},
},
},
newTree: TestTree{
"foo": TestFile{},
"subdir": TestTree{},
},
check: checkRewriteSkips(
map[string]struct{}{
"/subdir/subfile": {},
},
[]string{
"/foo",
"/subdir",
"/subdir/subfile",
},
false,
),
},
{ // exclude dir
tree: TestTree{
"foo": TestFile{},
"subdir": TestTree{
"subfile": TestFile{},
},
},
newTree: TestTree{
"foo": TestFile{},
},
check: checkRewriteSkips(
map[string]struct{}{
"/subdir": {},
},
[]string{
"/foo",
"/subdir",
},
false,
),
},
{ // modify node
tree: TestTree{
"foo": TestFile{Size: 21},
"subdir": TestTree{
"subfile": TestFile{Size: 21},
},
},
newTree: TestTree{
"foo": TestFile{Size: 42},
"subdir": TestTree{
"subfile": TestFile{Size: 42},
},
},
check: checkIncreaseNodeSize(21),
},
{ // test cache
tree: TestTree{
// both subdirs are identical
"subdir1": TestTree{
"subfile": TestFile{},
"subfile2": TestFile{},
},
"subdir2": TestTree{
"subfile": TestFile{},
"subfile2": TestFile{},
},
},
newTree: TestTree{
"subdir1": TestTree{
"subfile2": TestFile{},
},
"subdir2": TestTree{
"subfile2": TestFile{},
},
},
check: checkRewriteSkips(
map[string]struct{}{
"/subdir1/subfile": {},
},
[]string{
"/subdir1",
"/subdir1/subfile",
"/subdir1/subfile2",
"/subdir2",
},
false,
),
},
{ // test disabled cache
tree: TestTree{
// both subdirs are identical
"subdir1": TestTree{
"subfile": TestFile{},
"subfile2": TestFile{},
},
"subdir2": TestTree{
"subfile": TestFile{},
"subfile2": TestFile{},
},
},
newTree: TestTree{
"subdir1": TestTree{
"subfile2": TestFile{},
},
"subdir2": TestTree{
"subfile": TestFile{},
"subfile2": TestFile{},
},
},
check: checkRewriteSkips(
map[string]struct{}{
"/subdir1/subfile": {},
},
[]string{
"/subdir1",
"/subdir1/subfile",
"/subdir1/subfile2",
"/subdir2",
"/subdir2/subfile",
"/subdir2/subfile2",
},
true,
),
},
}
for _, test := range tests {
t.Run("", func(t *testing.T) {
repo, root := BuildTreeMap(test.tree)
if test.newTree == nil {
test.newTree = test.tree
}
expRepo, expRoot := BuildTreeMap(test.newTree)
modrepo := WritableTreeMap{repo}
ctx, cancel := context.WithCancel(context.TODO())
defer cancel()
rewriter, last := test.check(t)
newRoot, err := rewriter.RewriteTree(ctx, modrepo, "/", root)
if err != nil {
t.Error(err)
}
last(t)
// verifying against the expected tree root also implicitly checks the structural integrity
if newRoot != expRoot {
t.Error("hash mismatch")
fmt.Println("Got")
modrepo.Dump()
fmt.Println("Expected")
WritableTreeMap{expRepo}.Dump()
}
})
}
}
func TestSnapshotSizeQuery(t *testing.T) {
tree := TestTree{
"foo": TestFile{Size: 21},
"bar": TestFile{Size: 21},
"subdir": TestTree{
"subfile": TestFile{Size: 21},
},
}
newTree := TestTree{
"foo": TestFile{Size: 42},
"subdir": TestTree{
"subfile": TestFile{Size: 42},
},
}
t.Run("", func(t *testing.T) {
repo, root := BuildTreeMap(tree)
expRepo, expRoot := BuildTreeMap(newTree)
modrepo := WritableTreeMap{repo}
ctx, cancel := context.WithCancel(context.TODO())
defer cancel()
rewriteNode := func(node *restic.Node, path string) *restic.Node {
if path == "/bar" {
return nil
}
if node.Type == "file" {
node.Size += 21
}
return node
}
rewriter, querySize := NewSnapshotSizeRewriter(rewriteNode)
newRoot, err := rewriter.RewriteTree(ctx, modrepo, "/", root)
if err != nil {
t.Error(err)
}
ss := querySize()
test.Equals(t, uint(2), ss.FileCount, "snapshot file count mismatch")
test.Equals(t, uint64(84), ss.FileSize, "snapshot size mismatch")
// verifying against the expected tree root also implicitly checks the structural integrity
if newRoot != expRoot {
t.Error("hash mismatch")
fmt.Println("Got")
modrepo.Dump()
fmt.Println("Expected")
WritableTreeMap{expRepo}.Dump()
}
})
}
func TestRewriterFailOnUnknownFields(t *testing.T) {
tm := WritableTreeMap{TreeMap{}}
node := []byte(`{"nodes":[{"name":"subfile","type":"file","mtime":"0001-01-01T00:00:00Z","atime":"0001-01-01T00:00:00Z","ctime":"0001-01-01T00:00:00Z","uid":0,"gid":0,"content":null,"unknown_field":42}]}`)
id := restic.Hash(node)
tm.TreeMap[id] = node
ctx, cancel := context.WithCancel(context.TODO())
defer cancel()
rewriter := NewTreeRewriter(RewriteOpts{
RewriteNode: func(node *restic.Node, path string) *restic.Node {
// tree loading must not succeed
t.Fail()
return node
},
})
_, err := rewriter.RewriteTree(ctx, tm, "/", id)
if err == nil {
t.Error("missing error on unknown field")
}
// check that the serialization check can be disabled
rewriter = NewTreeRewriter(RewriteOpts{
AllowUnstableSerialization: true,
})
root, err := rewriter.RewriteTree(ctx, tm, "/", id)
test.OK(t, err)
_, expRoot := BuildTreeMap(TestTree{
"subfile": TestFile{},
})
test.Assert(t, root == expRoot, "mismatched trees")
}
func TestRewriterTreeLoadError(t *testing.T) {
tm := WritableTreeMap{TreeMap{}}
id := restic.NewRandomID()
ctx, cancel := context.WithCancel(context.TODO())
defer cancel()
// also check that load error by default cause the operation to fail
rewriter := NewTreeRewriter(RewriteOpts{})
_, err := rewriter.RewriteTree(ctx, tm, "/", id)
if err == nil {
t.Fatal("missing error on unloadable tree")
}
replacementID := restic.NewRandomID()
rewriter = NewTreeRewriter(RewriteOpts{
RewriteFailedTree: func(nodeID restic.ID, path string, err error) (restic.ID, error) {
if nodeID != id || path != "/" {
t.Fail()
}
return replacementID, nil
},
})
newRoot, err := rewriter.RewriteTree(ctx, tm, "/", id)
test.OK(t, err)
test.Equals(t, replacementID, newRoot)
}