From ef46b60e052e3e13a754754f7071b78ad33c6643 Mon Sep 17 00:00:00 2001
From: Minh Doan <minh@simple-rules.com>
Date: Mon, 10 Sep 2018 18:38:19 -0700
Subject: [PATCH] add test for gob encoding

---
 blockchain/blockchain.go      | 19 +++++++++++++++++++
 blockchain/blockchain_test.go |  3 ++-
 p2p/helper_test.go            | 33 ++++++++++++++++++++++++++++-----
 3 files changed, 49 insertions(+), 6 deletions(-)

diff --git a/blockchain/blockchain.go b/blockchain/blockchain.go
index f9c68a8aa..a549a42d2 100644
--- a/blockchain/blockchain.go
+++ b/blockchain/blockchain.go
@@ -15,6 +15,15 @@ type Blockchain struct {
 
 const genesisCoinbaseData = "The Times 03/Jan/2009 Chancellor on brink of second bailout for banks"
 
+func (bc *Blockchain) FindBlock(blockHash [32]byte) *Block {
+	for _, block := range bc.Blocks {
+		if bytes.Compare(block.Hash[:], blockHash[:]) == 0 {
+			return block
+		}
+	}
+	return nil
+}
+
 // GetLatestBlock gests the latest block at the end of the chain
 func (bc *Blockchain) GetLatestBlock() *Block {
 	if len(bc.Blocks) == 0 {
@@ -196,6 +205,16 @@ func CreateBlockchain(address [20]byte, shardId uint32) *Blockchain {
 	return &bc
 }
 
+// Creat testing genesis block.
+func CreateTestingGenesisBlock() *Block {
+	priIntOne := 111
+	shardId := uint32(1)
+	testAddressOne := pki.GetAddressFromInt(priIntOne)
+	cbtx := NewCoinbaseTX(testAddressOne, genesisCoinbaseData, shardId)
+	genesis := NewGenesisBlock(cbtx, shardId)
+	return genesis
+}
+
 // Create state block based on the utxos.
 func (bc *Blockchain) CreateStateBlock(utxoPool *UTXOPool) *Block {
 	var numBlocks int32 = 0
diff --git a/blockchain/blockchain_test.go b/blockchain/blockchain_test.go
index bccf9fee0..f83bc9d51 100644
--- a/blockchain/blockchain_test.go
+++ b/blockchain/blockchain_test.go
@@ -1,8 +1,9 @@
 package blockchain
 
 import (
-	"github.com/simple-rules/harmony-benchmark/crypto/pki"
 	"testing"
+
+	"github.com/simple-rules/harmony-benchmark/crypto/pki"
 )
 
 var (
diff --git a/p2p/helper_test.go b/p2p/helper_test.go
index 388749b5d..4f919549e 100644
--- a/p2p/helper_test.go
+++ b/p2p/helper_test.go
@@ -1,9 +1,15 @@
-package p2p
+package p2p_test
 
 import (
 	"bufio"
+	"bytes"
+	"encoding/gob"
 	"net"
+	"reflect"
 	"testing"
+
+	"github.com/simple-rules/harmony-benchmark/blockchain"
+	"github.com/simple-rules/harmony-benchmark/p2p"
 )
 
 func setUpTestServer(times int, t *testing.T, conCreated chan struct{}) {
@@ -18,11 +24,11 @@ func setUpTestServer(times int, t *testing.T, conCreated chan struct{}) {
 	)
 	for times > 0 {
 		times--
-		data, err := ReadMessageContent(conn)
+		data, err := p2p.ReadMessageContent(conn)
 		if err != nil {
 			t.Fatalf("error when ReadMessageContent %v", err)
 		}
-		data = CreateMessage(byte(1), data)
+		data = p2p.CreateMessage(byte(1), data)
 		w.Write(data)
 		w.Flush()
 	}
@@ -40,9 +46,9 @@ func TestNewNewNode(t *testing.T) {
 		times--
 
 		myMsg := "minhdoan"
-		SendMessageContent(conn, []byte(myMsg))
+		p2p.SendMessageContent(conn, []byte(myMsg))
 
-		data, err := ReadMessageContent(conn)
+		data, err := p2p.ReadMessageContent(conn)
 		if err != nil {
 			t.Error("got an error when trying to receive an expected message from server.")
 		}
@@ -51,3 +57,20 @@ func TestNewNewNode(t *testing.T) {
 		}
 	}
 }
+
+func TestGobEncode(t *testing.T) {
+	block := blockchain.CreateTestingGenesisBlock()
+
+	var tmp bytes.Buffer
+	enc := gob.NewEncoder(&tmp)
+	enc.Encode(*block)
+
+	tmp2 := bytes.NewBuffer(tmp.Bytes())
+	dec := gob.NewDecoder(tmp2)
+
+	var block2 blockchain.Block
+	dec.Decode(&block2)
+	if !reflect.DeepEqual(*block, block2) {
+		t.Error("Error in GobEncode")
+	}
+}