Add a test and fix merkle tree implementation

merkle_tree.go

@@ -20,20 +20,20 @@ type MerkleNode struct {
 func NewMerkleTree(data [][]byte) *MerkleTree {
 	var nodes []MerkleNode
 
-	if len(data)%2 != 0 {
-		data = append(data, data[len(data)-1])
-	}
-
 	for _, datum := range data {
 		node := NewMerkleNode(nil, nil, datum)
 		nodes = append(nodes, *node)
 	}
 
-	for i := 0; i < len(data)/2; i++ {
+  for len(nodes) > 1 {
 		var newLevel []MerkleNode
 
-		for j := 0; j < len(nodes); j += 2 {
-			node := NewMerkleNode(&nodes[j], &nodes[j+1], nil)
+		for i := 0; i < len(nodes); i += 2 {
+      j := i + 1
+      if j == len(nodes) {
+        j -= 1
+      }
+			node := NewMerkleNode(&nodes[i], &nodes[j], nil)
 			newLevel = append(newLevel, *node)
 		}
 

merkle_tree_test.go

@@ -73,3 +73,36 @@ func TestNewMerkleTree(t *testing.T) {
 
 	assert.Equal(t, rootHash, fmt.Sprintf("%x", mTree.RootNode.Data), "Merkle tree root hash is correct")
 }
+
+func TestNewMerkleTree2(t *testing.T) {
+	data := [][]byte{
+		[]byte("node1"),
+		[]byte("node2"),
+		[]byte("node3"),
+		[]byte("node4"),
+		[]byte("node5"),
+	}
+	// Level 1
+	n11 := NewMerkleNode(nil, nil, data[0])
+	n12 := NewMerkleNode(nil, nil, data[1])
+	n13 := NewMerkleNode(nil, nil, data[2])
+	n14 := NewMerkleNode(nil, nil, data[3])
+	n15 := NewMerkleNode(nil, nil, data[4])
+
+	// Level 2
+	n21 := NewMerkleNode(n11, n12, nil)
+	n22 := NewMerkleNode(n13, n14, nil)
+	n23 := NewMerkleNode(n15, n15, nil)
+
+	// Level 3
+	n31 := NewMerkleNode(n21, n22, nil)
+	n32 := NewMerkleNode(n23, n23, nil)
+
+	// Level 4
+	n41 := NewMerkleNode(n31, n32, nil)
+
+	rootHash := fmt.Sprintf("%x", n41.Data)
+	mTree := NewMerkleTree(data)
+
+	assert.Equal(t, rootHash, fmt.Sprintf("%x", mTree.RootNode.Data), "Merkle tree root hash is correct")
+}