use require package in unit tests

zeta-chain · Oct 18, 2023 · 26b1ab1 · 26b1ab1
1 parent f341067
commit 26b1ab1
Show file tree

Hide file tree

Showing 2 changed files with 26 additions and 56 deletions.
diff --git a/common/headers_test.go b/common/headers_test.go
@@ -4,7 +4,6 @@ import (
 	"bytes"
 	"encoding/base64"
 	"fmt"
-	"log"
 	"testing"
 	"time"
 
@@ -26,7 +25,7 @@ func TestTrueBitcoinHeader(t *testing.T) {
 		// Deserialize the header bytes from base64
 		headerBytes, err := base64.StdEncoding.DecodeString(b.HeaderBase64)
 		require.NoError(t, err)
-		header := unmarshalHeader(headerBytes)
+		header := unmarshalHeader(t, headerBytes)
 
 		// Validate
 		validateTrueBitcoinHeader(t, header, headerBytes)
@@ -40,7 +39,7 @@ func TestFakeBitcoinHeader(t *testing.T) {
 		// Deserialize the header bytes from base64
 		headerBytes, err := base64.StdEncoding.DecodeString(b.HeaderBase64)
 		require.NoError(t, err)
-		header := unmarshalHeader(headerBytes)
+		header := unmarshalHeader(t, headerBytes)
 
 		// Validate
 		validateFakeBitcoinHeader(t, header, headerBytes)
@@ -60,7 +59,7 @@ func BitcoinHeaderValidationLiveTest(t *testing.T) {
 		// Get the block header
 		header, err := client.GetBlockHeader(blockHash)
 		require.NoError(t, err)
-		headerBytes := marshalHeader(header)
+		headerBytes := marshalHeader(t, header)
 
 		// Validate true header
 		validateTrueBitcoinHeader(t, header, headerBytes)
@@ -101,21 +100,17 @@ func copyHeader(header *wire.BlockHeader) *wire.BlockHeader {
 	return copyHeader
 }
 
-func marshalHeader(header *wire.BlockHeader) []byte {
+func marshalHeader(t *testing.T, header *wire.BlockHeader) []byte {
 	var headerBuf bytes.Buffer
 	err := header.Serialize(&headerBuf)
-	if err != nil {
-		log.Fatal(err)
-	}
+	require.NoError(t, err)
 	return headerBuf.Bytes()
 }
 
-func unmarshalHeader(headerBytes []byte) *wire.BlockHeader {
+func unmarshalHeader(t *testing.T, headerBytes []byte) *wire.BlockHeader {
 	var header wire.BlockHeader
 	err := header.Deserialize(bytes.NewReader(headerBytes))
-	if err != nil {
-		log.Fatal(err)
-	}
+	require.NoError(t, err)
 	return &header
 }
 
@@ -136,52 +131,40 @@ func validateFakeBitcoinHeader(t *testing.T, header *wire.BlockHeader, headerByt
 
 	// Incorrect header length should fail validation
 	err := common.ValidateBitcoinHeader(headerBytes[:79], blockHash[:], 18332)
-	if err == nil {
-		t.Error("Incorrect header length should fail validation")
-	}
+	require.Error(t, err)
 
 	// Incorrect version should fail validation
 	fakeHeader := copyHeader(header)
 	fakeHeader.Version = 0
-	fakeBytes := marshalHeader(fakeHeader)
+	fakeBytes := marshalHeader(t, fakeHeader)
 	fakeHash := fakeHeader.BlockHash()
 	err = common.ValidateBitcoinHeader(fakeBytes, fakeHash[:], 18332)
-	if err == nil {
-		t.Error("Incorrect version should fail validation")
-	}
+	require.Error(t, err)
 
 	// Incorrect timestamp should fail validation
 	// Case1: timestamp is before genesis block
 	fakeHeader = copyHeader(header)
 	fakeHeader.Timestamp = chaincfg.TestNet3Params.GenesisBlock.Header.Timestamp.Add(-time.Second)
-	fakeBytes = marshalHeader(fakeHeader)
+	fakeBytes = marshalHeader(t, fakeHeader)
 	fakeHash = fakeHeader.BlockHash()
 	err = common.ValidateBitcoinHeader(fakeBytes, fakeHash[:], 18332)
-	if err == nil {
-		t.Error("Timestamp before genesis should fail validation")
-	}
+	require.Error(t, err)
 	// Case2: timestamp is after 2 hours in the future
 	fakeHeader = copyHeader(header)
 	fakeHeader.Timestamp = header.Timestamp.Add(time.Second * (blockchain.MaxTimeOffsetSeconds + 1))
-	fakeBytes = marshalHeader(fakeHeader)
+	fakeBytes = marshalHeader(t, fakeHeader)
 	err = common.NewBitcoinHeader(fakeBytes).ValidateTimestamp(header.Timestamp)
-	if err == nil {
-		t.Error("Timestamp in future should fail validation")
-	}
+	require.Error(t, err)
 
 	// Incorrect block hash should fail validation
 	fakeHeader = copyHeader(header)
 	header.Nonce = 0
-	fakeBytes = marshalHeader(header)
+	fakeBytes = marshalHeader(t, header)
 	err = common.ValidateBitcoinHeader(fakeBytes, blockHash[:], 18332)
-	if err == nil {
-		t.Error("Incorrect block hash should fail validation")
-	}
+	require.Error(t, err)
 
 	// PoW not satisfied should fail validation
 	fakeHash = fakeHeader.BlockHash()
 	err = common.ValidateBitcoinHeader(fakeBytes, fakeHash[:], 18332)
-	if err == nil {
-		t.Error("PoW not satisfied should fail validation")
-	}
+	require.Error(t, err)
 }
diff --git a/common/proof_test.go b/common/proof_test.go
@@ -5,12 +5,10 @@ import (
 	"os"
 	"testing"
 
-	"bytes"
 	"encoding/base64"
 	"encoding/hex"
 	"encoding/json"
 	"fmt"
-	"log"
 
 	"github.com/stretchr/testify/require"
 	"github.com/zeta-chain/zetacore/common"
@@ -65,7 +63,7 @@ func TestBitcoinMerkleProof(t *testing.T) {
 		// Deserialize the header bytes from base64
 		headerBytes, err := base64.StdEncoding.DecodeString(b.HeaderBase64)
 		require.NoError(t, err)
-		header := unmarshalHeader(headerBytes)
+		header := unmarshalHeader(t, headerBytes)
 
 		// Deserialize the block bytes from base64
 		blockBytes, err := base64.StdEncoding.DecodeString(b.BlockBase64)
@@ -93,7 +91,7 @@ func BitcoinMerkleProofLiveTest(t *testing.T) {
 		// Get the block header
 		header, err := client.GetBlockHeader(blockHash)
 		require.NoError(t, err)
-		headerBytes := marshalHeader(header)
+		headerBytes := marshalHeader(t, header)
 		target := blockchain.CompactToBig(header.Bits)
 
 		// Get the block with verbose transactions
@@ -113,13 +111,9 @@ func validateBitcoinBlock(t *testing.T, header *wire.BlockHeader, headerBytes []
 	txBodies := [][]byte{}
 	for _, res := range blockVerbose.Tx {
 		txBytes, err := hex.DecodeString(res.Hex)
-		if err != nil {
-			log.Fatalf("error decoding transaction hex: %v", err)
-		}
+		require.NoError(t, err)
 		tx, err := btcutil.NewTxFromBytes(txBytes)
-		if err != nil {
-			log.Fatalf("error deserializing transaction: %v", err)
-		}
+		require.NoError(t, err)
 
 		// Validate Tss SegWit transaction if it's an outTx
 		if res.Txid == outTxid {
@@ -139,26 +133,19 @@ func validateBitcoinBlock(t *testing.T, header *wire.BlockHeader, headerBytes []
 	mk := bitcoin.NewMerkle(txns)
 	for i := range txns {
 		path, index, err := mk.BuildMerkleProof(i)
-		if err != nil {
-			log.Fatalf("Error building merkle proof: %v", err)
-		}
+		require.NoError(t, err)
 
 		// True proof should verify
 		proof := common.NewBitcoinProof(txBodies[i], path, index)
 		txBytes, err := proof.Verify(common.NewBitcoinHeader(headerBytes), 0)
-		if err != nil {
-			log.Fatal("Merkle proof verification failed")
-		}
-		if !bytes.Equal(txBytes, txBodies[i]) {
-			log.Fatalf("Transaction body mismatch")
-		}
+		require.NoError(t, err)
+		require.Equal(t, txBytes, txBodies[i])
 
 		// Fake proof should not verify
 		fakeIndex := index ^ 0xffffffff // flip all bits
 		fakeProof := common.NewBitcoinProof(txBodies[i], path, fakeIndex)
 		txBytes, err = fakeProof.Verify(common.NewBitcoinHeader(headerBytes), 0)
-		if err == nil || txBytes != nil {
-			log.Fatalf("Merkle proof should not verify")
-		}
+		require.Error(t, err)
+		require.Nil(t, txBytes)
 	}
 }