From f6c514f802effbc7dd9955e38abd45c7e0ea91b8 Mon Sep 17 00:00:00 2001
From: Minh Doan <minh@simple-rules.com>
Date: Sun, 25 Nov 2018 23:29:34 -0800
Subject: [PATCH] add test for GetBlocks

---
 syncing/downloader/client.go      | 14 ++++++++++
 syncing/downloader/server_test.go | 43 +++++++++++++++++++++++++++----
 2 files changed, 52 insertions(+), 5 deletions(-)

diff --git a/syncing/downloader/client.go b/syncing/downloader/client.go
index d4418182c..976452af4 100644
--- a/syncing/downloader/client.go
+++ b/syncing/downloader/client.go
@@ -59,3 +59,17 @@ func (client *Client) GetBlockHashes() *pb.DownloaderResponse {
 	}
 	return response
 }
+
+// GetBlocks ...
+func (client *Client) GetBlocks(heights []int32) *pb.DownloaderResponse {
+	ctx, cancel := context.WithTimeout(context.Background(), 10*time.Second)
+	defer cancel()
+	request := &pb.DownloaderRequest{Type: pb.DownloaderRequest_BLOCK}
+	request.Height = make([]int32, len(heights))
+	copy(request.Height, heights)
+	response, err := client.dlClient.Query(ctx, request)
+	if err != nil {
+		log.Fatalf("Error")
+	}
+	return response
+}
diff --git a/syncing/downloader/server_test.go b/syncing/downloader/server_test.go
index 0e80af69f..098abe9c8 100644
--- a/syncing/downloader/server_test.go
+++ b/syncing/downloader/server_test.go
@@ -1,4 +1,4 @@
-package downloader
+package downloader_test
 
 import (
 	"fmt"
@@ -7,6 +7,7 @@ import (
 
 	bc "github.com/harmony-one/harmony/blockchain"
 	"github.com/harmony-one/harmony/crypto/pki"
+	"github.com/harmony-one/harmony/syncing/downloader"
 	pb "github.com/harmony-one/harmony/syncing/downloader/proto"
 )
 
@@ -37,6 +38,15 @@ func (node *FakeNode) GetBlockHashes() [][]byte {
 	return res
 }
 
+// GetBlocks used for state download.
+func (node *FakeNode) GetBlocks() [][]byte {
+	res := [][]byte{}
+	for _, block := range node.bc.Blocks {
+		res = append(res, block.Serialize())
+	}
+	return res
+}
+
 // SetBlockchain is used for testing
 func (node *FakeNode) Init() {
 	addresses := [][20]byte{TestAddressOne, TestAddressTwo}
@@ -46,12 +56,13 @@ func (node *FakeNode) Init() {
 func (node *FakeNode) CalculateResponse(request *pb.DownloaderRequest) (*pb.DownloaderResponse, error) {
 	response := &pb.DownloaderResponse{}
 	if request.Type == pb.DownloaderRequest_HEADER {
-		fmt.Println("minh ", len(node.bc.Blocks))
 		for _, block := range node.bc.Blocks {
 			response.Payload = append(response.Payload, block.Hash[:])
 		}
 	} else {
-
+		for _, id := range request.Height {
+			response.Payload = append(response.Payload, node.bc.Blocks[id].Serialize())
+		}
 	}
 	return response, nil
 }
@@ -59,18 +70,40 @@ func (node *FakeNode) CalculateResponse(request *pb.DownloaderRequest) (*pb.Down
 func TestGetBlockHashes(t *testing.T) {
 	fakeNode := &FakeNode{}
 	fakeNode.Init()
-	s := NewServer(fakeNode)
+	s := downloader.NewServer(fakeNode)
 	grcpServer, err := s.Start(serverIP, serverPort)
 	if err != nil {
 		t.Error(err)
 	}
 	defer grcpServer.Stop()
 
-	client := ClientSetup(serverIP, serverPort)
+	client := downloader.ClientSetup(serverIP, serverPort)
+	defer client.Close()
 	response := client.GetBlockHashes()
 	if !reflect.DeepEqual(response.Payload, fakeNode.GetBlockHashes()) {
 		t.Error("not equal")
 	}
+}
+
+func TestGetBlocks(t *testing.T) {
+	fakeNode := &FakeNode{}
+	fakeNode.Init()
+	s := downloader.NewServer(fakeNode)
+	grcpServer, err := s.Start(serverIP, serverPort)
+	if err != nil {
+		t.Error(err)
+	}
+	defer grcpServer.Stop()
 
+	client := downloader.ClientSetup(serverIP, serverPort)
 	defer client.Close()
+	response := client.GetBlockHashes()
+	if !reflect.DeepEqual(response.Payload, fakeNode.GetBlockHashes()) {
+		t.Error("not equal")
+	}
+	response = client.GetBlocks([]int32{0, 1})
+	fmt.Println(len(response.Payload))
+	if !reflect.DeepEqual(response.Payload, fakeNode.GetBlocks()) {
+		t.Error("not equal")
+	}
 }