From e7204d29b1e83126af12ad1ebdee20c1ace4f426 Mon Sep 17 00:00:00 2001
From: txaty <txaty@proton.me>
Date: Fri, 24 Nov 2023 01:42:23 +0800
Subject: [PATCH] Fix unit test issues.

Signed-off-by: txaty <txaty@proton.me>
---
 merkle_tree_test.go              |  2 +-
 proof_gen_and_tree_build_test.go |  7 +++++--
 tree_build_test.go               | 19 +++++++++++++++++++
 3 files changed, 25 insertions(+), 3 deletions(-)

diff --git a/merkle_tree_test.go b/merkle_tree_test.go
index cf12f8e..80bcf1b 100644
--- a/merkle_tree_test.go
+++ b/merkle_tree_test.go
@@ -34,7 +34,7 @@ import (
 func mockDataBlocks(num int) []DataBlock {
 	blocks := make([]DataBlock, num)
 	for i := 0; i < num; i++ {
-		byteLen := rand.Intn(1 << 15)
+		byteLen := max(32, rand.Intn(1<<15))
 		block := &mock.DataBlock{
 			Data: make([]byte, byteLen),
 		}
diff --git a/proof_gen_and_tree_build_test.go b/proof_gen_and_tree_build_test.go
index 6e7ea33..9d05df0 100644
--- a/proof_gen_and_tree_build_test.go
+++ b/proof_gen_and_tree_build_test.go
@@ -26,6 +26,7 @@ import (
 	"crypto/sha256"
 	"fmt"
 	"reflect"
+	"sync/atomic"
 	"testing"
 )
 
@@ -147,6 +148,7 @@ func TestMerkleTreeNew_modeProofGenAndTreeBuild(t *testing.T) {
 }
 
 func TestMerkleTreeNew_modeProofGenAndTreeBuildParallel(t *testing.T) {
+	var hashFuncCounter atomic.Uint32
 	type args struct {
 		blocks []DataBlock
 		config *Config
@@ -221,12 +223,13 @@ func TestMerkleTreeNew_modeProofGenAndTreeBuildParallel(t *testing.T) {
 		{
 			name: "test_tree_build_hash_func_error",
 			args: args{
-				blocks: mockDataBlocks(100),
+				blocks: mockDataBlocks(5),
 				config: &Config{
 					HashFunc: func(block []byte) ([]byte, error) {
-						if len(block) == 64 {
+						if hashFuncCounter.Load() >= 5 {
 							return nil, fmt.Errorf("hash func error")
 						}
+						hashFuncCounter.Add(1)
 						sha256Func := sha256.New()
 						sha256Func.Write(block)
 						return sha256Func.Sum(nil), nil
diff --git a/tree_build_test.go b/tree_build_test.go
index f076754..5e4816f 100644
--- a/tree_build_test.go
+++ b/tree_build_test.go
@@ -276,6 +276,24 @@ func TestMerkleTreeNew_modeTreeBuildParallel(t *testing.T) {
 			},
 			wantErr: true,
 		},
+		{
+			name: "test_hash_func_error_when_computing_nodes_parallel",
+			args: args{
+				blocks: mockDataBlocks(4),
+				config: &Config{
+					HashFunc: func(block []byte) ([]byte, error) {
+						if hashFuncCounter.Load() == 5 {
+							return nil, fmt.Errorf("hash func error")
+						}
+						hashFuncCounter.Add(1)
+						sha256Func := sha256.New()
+						sha256Func.Write(block)
+						return sha256Func.Sum(nil), nil
+					},
+				},
+			},
+			wantErr: true,
+		},
 		{
 			name: "test_hash_func_error_when_computing_root_parallel",
 			args: args{
@@ -299,6 +317,7 @@ func TestMerkleTreeNew_modeTreeBuildParallel(t *testing.T) {
 	}
 	for _, tt := range tests {
 		t.Run(tt.name, func(t *testing.T) {
+			hashFuncCounter.Store(0)
 			m, err := New(tt.args.config, tt.args.blocks)
 			if (err != nil) != tt.wantErr {
 				t.Errorf("Build() error = %v, wantErr %v", err, tt.wantErr)