correct nested diff colors and dir sizing (fixes #45)

4 files changed, 36 insertions, 12 deletions

diff --git a/filetree/node.go b/filetree/node.go
index 95cc854..233692a 100644
--- a/filetree/node.go
+++ b/filetree/node.go
@@ -147,18 +147,19 @@ func (node *FileNode) MetadataString() string {
 
 	var sizeBytes int64
 
-	if node.Data.FileInfo.TarHeader.FileInfo().IsDir() {
-
+	if node.IsLeaf() {
+		sizeBytes = node.Data.FileInfo.TarHeader.FileInfo().Size()
+	} else {
 		sizer := func(curNode *FileNode) error {
-			if curNode.Data.DiffType != Removed {
+			// don't include file sizes of children that have been removed (unless the node in question is a removed dir,
+			// then show the accumulated size of removed files)
+			if curNode.Data.DiffType != Removed || node.Data.DiffType == Removed {
 				sizeBytes += curNode.Data.FileInfo.TarHeader.FileInfo().Size()
 			}
 			return nil
 		}
 
 		node.VisitDepthChildFirst(sizer, nil)
-	} else {
-		sizeBytes = node.Data.FileInfo.TarHeader.FileInfo().Size()
 	}
 
 	size := humanize.Bytes(uint64(sizeBytes))
@@ -254,7 +255,7 @@ func (node *FileNode) Path() string {
 		}
 		node.path = "/" + strings.Join(path, "/")
 	}
-	return node.path
+	return strings.Replace(node.path, "//", "/", -1)
 }
 
 // deriveDiffType determines a DiffType to the current FileNode. Note: the DiffType of a node is always the DiffType of
diff --git a/filetree/node_test.go b/filetree/node_test.go
index 1111444..db39d7d 100644
--- a/filetree/node_test.go
+++ b/filetree/node_test.go
@@ -1,6 +1,7 @@
 package filetree
 
 import (
+	"archive/tar"
 	"testing"
 )
 
@@ -87,7 +88,7 @@ func TestPath(t *testing.T) {
 
 	actual := node.Path()
 	if expected != actual {
-		t.Errorf("Expected Path '%s' got '%s'", expected, actual)
+		t.Errorf("Expected path '%s' got '%s'", expected, actual)
 	}
 }
 
@@ -148,3 +149,16 @@ func TestDiffTypeFromRemovedChildren(t *testing.T) {
 	}
 
 }
+
+func TestDirSize(t *testing.T) {
+	tree1 := NewFileTree()
+	tree1.AddPath("/etc/nginx/public1", FileInfo{TarHeader: tar.Header{Size: 100}})
+	tree1.AddPath("/etc/nginx/thing1", FileInfo{TarHeader: tar.Header{Size: 200}})
+	tree1.AddPath("/etc/nginx/public3/thing2", FileInfo{TarHeader: tar.Header{Size: 300}})
+
+	node, _ := tree1.GetNode("/etc/nginx")
+	expected, actual := "----------        0:0      600 B ", node.MetadataString()
+	if expected != actual {
+		t.Errorf("Expected metadata '%s' got '%s'", expected, actual)
+	}
+}
diff --git a/filetree/tree.go b/filetree/tree.go
index 41df47a..d87e7ec 100644
--- a/filetree/tree.go
+++ b/filetree/tree.go
@@ -245,8 +245,11 @@ func (tree *FileTree) RemovePath(path string) error {
 	return node.Remove()
 }
 
-// Compare marks the FileNodes in the owning tree with DiffType annotations when compared to the given tree.
+// Compare marks the FileNodes in the owning (lower) tree with DiffType annotations when compared to the given (upper) tree.
 func (tree *FileTree) Compare(upper *FileTree) error {
+	// always compare relative to the original, unaltered tree.
+	originalTree := tree.Copy()
+
 	graft := func(upperNode *FileNode) error {
 		if upperNode.IsWhiteout() {
 			err := tree.markRemoved(upperNode.Path())
@@ -254,20 +257,26 @@ func (tree *FileTree) Compare(upper *FileTree) error {
 				return fmt.Errorf("cannot remove upperNode %s: %v", upperNode.Path(), err.Error())
 			}
 		} else {
-			lowerNode, _ := tree.GetNode(upperNode.Path())
-			if lowerNode == nil {
+			// compare against the original tree (to ensure new parents with new children are captured as new instead of modified)
+			originalLowerNode, _ := originalTree.GetNode(upperNode.Path())
+
+			if originalLowerNode == nil {
 				newNode, err := tree.AddPath(upperNode.Path(), upperNode.Data.FileInfo)
 				if err != nil {
 					return fmt.Errorf("cannot add new upperNode %s: %v", upperNode.Path(), err.Error())
 				}
 				newNode.AssignDiffType(Added)
 			} else {
+				// check the tree for comparison markings
+				lowerNode, _ := tree.GetNode(upperNode.Path())
+
 				diffType := lowerNode.compare(upperNode)
 				return lowerNode.deriveDiffType(diffType)
 			}
 		}
 		return nil
 	}
+	// we must visit from the leaves upwards to ensure that diff types can be derived from and assigned to children
 	return upper.VisitDepthChildFirst(graft, nil)
 }
 
diff --git a/filetree/tree_test.go b/filetree/tree_test.go
index 186d19e..04f8b37 100644
--- a/filetree/tree_test.go
+++ b/filetree/tree_test.go
@@ -288,7 +288,7 @@ func TestCompareWithAdds(t *testing.T) {
 	lowerTree := NewFileTree()
 	upperTree := NewFileTree()
 	lowerPaths := [...]string{"/etc", "/etc/sudoers", "/usr", "/etc/hosts", "/usr/bin"}
-	upperPaths := [...]string{"/etc", "/etc/sudoers", "/usr", "/etc/hosts", "/usr/bin", "/usr/bin/bash"}
+	upperPaths := [...]string{"/etc", "/etc/sudoers", "/usr", "/etc/hosts", "/usr/bin", "/usr/bin/bash", "/a/new/path"}
 
 	for _, value := range lowerPaths {
 		lowerTree.AddPath(value, FileInfo{
@@ -316,7 +316,7 @@ func TestCompareWithAdds(t *testing.T) {
 		p := n.Path()
 		if p == "/" {
 			return nil
-		} else if stringInSlice(p, []string{"/usr/bin/bash"}) {
+		} else if stringInSlice(p, []string{"/usr/bin/bash", "/a", "/a/new", "/a/new/path"}) {
 			if err := AssertDiffType(n, Added); err != nil {
 				failedAssertions = append(failedAssertions, err)
 			}