restorer: Add tests for traverseTree

internal/restorer/restorer_test.go

@@ -466,3 +466,213 @@ func TestRestorerRelative(t *testing.T) {
 		})
 	}
 }
+
+type TraverseTreeCheck func(testing.TB) treeVisitor
+
+type TreeVisit struct {
+	funcName string // name of the function
+	location string // location passed to the function
+}
+
+func checkVisitOrder(list []TreeVisit) TraverseTreeCheck {
+	var pos int
+
+	return func(t testing.TB) treeVisitor {
+		check := func(funcName string) func(*restic.Node, string, string) error {
+			return func(node *restic.Node, target, location string) error {
+				if pos >= len(list) {
+					t.Errorf("step %v, %v(%v): expected no more than %d function calls", pos, funcName, location, len(list))
+					pos++
+					return nil
+				}
+
+				v := list[pos]
+
+				if v.funcName != funcName {
+					t.Errorf("step %v, location %v: want function %v, but %v was called",
+						pos, location, v.funcName, funcName)
+				}
+
+				if location != filepath.FromSlash(v.location) {
+					t.Errorf("step %v: want location %v, got %v", pos, list[pos].location, location)
+				}
+
+				pos++
+				return nil
+			}
+		}
+
+		return treeVisitor{
+			enterDir:  check("enterDir"),
+			visitNode: check("visitNode"),
+			leaveDir:  check("leaveDir"),
+		}
+	}
+}
+
+func TestRestorerTraverseTree(t *testing.T) {
+	var tests = []struct {
+		Snapshot
+		Select  func(item string, dstpath string, node *restic.Node) (selectForRestore bool, childMayBeSelected bool)
+		Visitor TraverseTreeCheck
+	}{
+		{
+			// select everything
+			Snapshot: Snapshot{
+				Nodes: map[string]Node{
+					"dir": Dir{Nodes: map[string]Node{
+						"otherfile": File{"x"},
+						"subdir": Dir{Nodes: map[string]Node{
+							"file": File{"content: file\n"},
+						}},
+					}},
+					"foo": File{"content: foo\n"},
+				},
+			},
+			Select: func(item string, dstpath string, node *restic.Node) (selectForRestore bool, childMayBeSelected bool) {
+				return true, true
+			},
+			Visitor: checkVisitOrder([]TreeVisit{
+				{"enterDir", "/dir"},
+				{"visitNode", "/dir/otherfile"},
+				{"enterDir", "/dir/subdir"},
+				{"visitNode", "/dir/subdir/file"},
+				{"leaveDir", "/dir/subdir"},
+				{"leaveDir", "/dir"},
+				{"visitNode", "/foo"},
+			}),
+		},
+
+		// select only the top-level file
+		{
+			Snapshot: Snapshot{
+				Nodes: map[string]Node{
+					"dir": Dir{Nodes: map[string]Node{
+						"otherfile": File{"x"},
+						"subdir": Dir{Nodes: map[string]Node{
+							"file": File{"content: file\n"},
+						}},
+					}},
+					"foo": File{"content: foo\n"},
+				},
+			},
+			Select: func(item string, dstpath string, node *restic.Node) (selectForRestore bool, childMayBeSelected bool) {
+				if item == "/foo" {
+					return true, false
+				}
+				return false, false
+			},
+			Visitor: checkVisitOrder([]TreeVisit{
+				{"visitNode", "/foo"},
+			}),
+		},
+		{
+			Snapshot: Snapshot{
+				Nodes: map[string]Node{
+					"aaa": File{"content: foo\n"},
+					"dir": Dir{Nodes: map[string]Node{
+						"otherfile": File{"x"},
+						"subdir": Dir{Nodes: map[string]Node{
+							"file": File{"content: file\n"},
+						}},
+					}},
+				},
+			},
+			Select: func(item string, dstpath string, node *restic.Node) (selectForRestore bool, childMayBeSelected bool) {
+				if item == "/aaa" {
+					return true, false
+				}
+				return false, false
+			},
+			Visitor: checkVisitOrder([]TreeVisit{
+				{"visitNode", "/aaa"},
+			}),
+		},
+
+		// select dir/
+		{
+			Snapshot: Snapshot{
+				Nodes: map[string]Node{
+					"dir": Dir{Nodes: map[string]Node{
+						"otherfile": File{"x"},
+						"subdir": Dir{Nodes: map[string]Node{
+							"file": File{"content: file\n"},
+						}},
+					}},
+					"foo": File{"content: foo\n"},
+				},
+			},
+			Select: func(item string, dstpath string, node *restic.Node) (selectForRestore bool, childMayBeSelected bool) {
+				if strings.HasPrefix(item, "/dir") {
+					return true, true
+				}
+				return false, false
+			},
+			Visitor: checkVisitOrder([]TreeVisit{
+				{"enterDir", "/dir"},
+				{"visitNode", "/dir/otherfile"},
+				{"enterDir", "/dir/subdir"},
+				{"visitNode", "/dir/subdir/file"},
+				{"leaveDir", "/dir/subdir"},
+				{"leaveDir", "/dir"},
+			}),
+		},
+
+		// select only dir/otherfile
+		{
+			Snapshot: Snapshot{
+				Nodes: map[string]Node{
+					"dir": Dir{Nodes: map[string]Node{
+						"otherfile": File{"x"},
+						"subdir": Dir{Nodes: map[string]Node{
+							"file": File{"content: file\n"},
+						}},
+					}},
+					"foo": File{"content: foo\n"},
+				},
+			},
+			Select: func(item string, dstpath string, node *restic.Node) (selectForRestore bool, childMayBeSelected bool) {
+				switch item {
+				case "/dir":
+					return false, true
+				case "/dir/otherfile":
+					return true, false
+				default:
+					return false, false
+				}
+			},
+			Visitor: checkVisitOrder([]TreeVisit{
+				{"visitNode", "/dir/otherfile"},
+			}),
+		},
+	}
+
+	for _, test := range tests {
+		t.Run("", func(t *testing.T) {
+			repo, cleanup := repository.TestRepository(t)
+			defer cleanup()
+			sn, id := saveSnapshot(t, repo, test.Snapshot)
+
+			res, err := NewRestorer(repo, id)
+			if err != nil {
+				t.Fatal(err)
+			}
+
+			res.SelectFilter = test.Select
+
+			tempdir, cleanup := rtest.TempDir(t)
+			defer cleanup()
+
+			ctx, cancel := context.WithCancel(context.Background())
+			defer cancel()
+
+			// make sure we're creating a new subdir of the tempdir
+			target := filepath.Join(tempdir, "target")
+
+			err = res.traverseTree(ctx, target, string(filepath.Separator), *sn.Tree, test.Visitor(t))
+			if err != nil {
+				t.Fatal(err)
+			}
+		})
+	}
+}