add test for merkle tree

blockchain/merkle_tree.go

@@ -18,6 +18,9 @@ type MerkleNode struct {
 
 // NewMerkleTree creates a new Merkle tree from a sequence of data
 func NewMerkleTree(data [][]byte) *MerkleTree {
+	if len(data) == 0 {
+		return nil
+	}
 	var nodes []*MerkleNode
 
 	for _, datum := range data {
@@ -49,15 +52,16 @@ func NewMerkleTree(data [][]byte) *MerkleTree {
 func NewMerkleNode(left, right *MerkleNode, data []byte) *MerkleNode {
 	mNode := MerkleNode{}
 
-	if left == nil && right == nil {
-		hash := sha256.Sum256(data)
-		mNode.Data = hash[:]
-	} else {
-		prevHashes := append(left.Data, right.Data...)
+	prevHashes := []byte{}
+	if left != nil {
+		prevHashes = append(prevHashes, left.Data...)
+	}
+	if right != nil {
+		prevHashes = append(prevHashes, right.Data...)
+	}
+	prevHashes = append(prevHashes, data...)
 	hash := sha256.Sum256(prevHashes)
 	mNode.Data = hash[:]
-	}
-
 	mNode.Left = left
 	mNode.Right = right
 

blockchain/merkle_tree_test.go

@@ -13,6 +13,7 @@ func TestNewMerkleNode(t *testing.T) {
 		[]byte("node3"),
 	}
 
+	fmt.Println("TEting")
 	// Level 1
 
 	n1 := NewMerkleNode(nil, nil, data[0])
@@ -59,3 +60,23 @@ func TestNewMerkleTree(t *testing.T) {
 		t.Errorf("Merkle tree root hash is incorrect")
 	}
 }
+
+func TestNewMerkleTree2(t *testing.T) {
+	data := [][]byte{
+		[]byte("node1"),
+		[]byte("node2"),
+	}
+	// Level 1
+	n1 := NewMerkleNode(nil, nil, data[0])
+	n2 := NewMerkleNode(nil, nil, data[1])
+
+	// Level 2
+	n3 := NewMerkleNode(n1, n2, nil)
+
+	rootHash := fmt.Sprintf("%x", n3.Data)
+	mTree := NewMerkleTree(data)
+
+	if rootHash != fmt.Sprintf("%x", mTree.RootNode.Data) {
+		t.Errorf("Merkle tree root hash is incorrect")
+	}
+}