dex: fix tests

4 files changed, 93 insertions, 43 deletions

diff --git a/dex/cache_test.go b/dex/cache_test.go
index 0ab426f3a..3c04db673 100644
--- a/dex/cache_test.go
+++ b/dex/cache_test.go
@@ -18,12 +18,28 @@
 package dex
 
 import (
+	"sort"
+	"strings"
 	"testing"
 
 	coreCommon "github.com/dexon-foundation/dexon-consensus/common"
 	coreTypes "github.com/dexon-foundation/dexon-consensus/core/types"
 )
 
+type byHash []*coreTypes.Vote
+
+func (v byHash) Len() int {
+	return len(v)
+}
+
+func (v byHash) Less(i int, j int) bool {
+	return strings.Compare(v[i].BlockHash.String(), v[j].BlockHash.String()) < 0
+}
+
+func (v byHash) Swap(i int, j int) {
+	v[i], v[j] = v[j], v[i]
+}
+
 func TestCacheVote(t *testing.T) {
 	cache := newCache(3)
 	pos0 := coreTypes.Position{
@@ -33,36 +49,50 @@ func TestCacheVote(t *testing.T) {
 		Height: uint64(1),
 	}
 	vote1 := &coreTypes.Vote{
-		BlockHash: coreCommon.NewRandomHash(),
-		Position:  pos0,
+		VoteHeader: coreTypes.VoteHeader{
+			BlockHash: coreCommon.NewRandomHash(),
+			Position:  pos0,
+		},
 	}
 	vote2 := &coreTypes.Vote{
-		BlockHash: coreCommon.NewRandomHash(),
-		Position:  pos0,
+		VoteHeader: coreTypes.VoteHeader{
+			BlockHash: coreCommon.NewRandomHash(),
+			Position:  pos0,
+		},
 	}
 	vote3 := &coreTypes.Vote{
-		BlockHash: coreCommon.NewRandomHash(),
-		Position:  pos1,
+		VoteHeader: coreTypes.VoteHeader{
+			BlockHash: coreCommon.NewRandomHash(),
+			Position:  pos1,
+		},
 	}
 	vote4 := &coreTypes.Vote{
-		BlockHash: coreCommon.NewRandomHash(),
-		Position:  pos1,
+		VoteHeader: coreTypes.VoteHeader{
+			BlockHash: coreCommon.NewRandomHash(),
+			Position:  pos1,
+		},
 	}
 	cache.addVote(vote1)
 	cache.addVote(vote2)
 	cache.addVote(vote3)
 
 	votes := cache.votes(pos0)
+	sort.Sort(byHash(votes))
+
+	resultVotes := []*coreTypes.Vote{vote1, vote2}
+	sort.Sort(byHash(resultVotes))
+
 	if len(votes) != 2 {
 		t.Errorf("fail to get votes: have %d, want 2", len(votes))
 	}
-	if !votes[0].BlockHash.Equal(vote1.BlockHash) {
-		t.Errorf("get wrong vote: have %s, want %s", votes[0], vote1)
+	if !votes[0].BlockHash.Equal(resultVotes[0].BlockHash) {
+		t.Errorf("get wrong vote: have %s, want %s", votes[0], resultVotes[0])
 	}
-	if !votes[1].BlockHash.Equal(vote2.BlockHash) {
-		t.Errorf("get wrong vote: have %s, want %s", votes[1], vote2)
+	if !votes[1].BlockHash.Equal(resultVotes[1].BlockHash) {
+		t.Errorf("get wrong vote: have %s, want %s", votes[1], resultVotes[1])
 	}
 	votes = cache.votes(pos1)
+	sort.Sort(byHash(votes))
 	if len(votes) != 1 {
 		t.Errorf("fail to get votes: have %d, want 1", len(votes))
 	}
@@ -73,18 +103,25 @@ func TestCacheVote(t *testing.T) {
 	cache.addVote(vote4)
 
 	votes = cache.votes(pos0)
+	sort.Sort(byHash(votes))
+
 	if len(votes) != 0 {
 		t.Errorf("fail to get votes: have %d, want 0", len(votes))
 	}
 	votes = cache.votes(pos1)
+	sort.Sort(byHash(votes))
+
+	resultVotes = []*coreTypes.Vote{vote3, vote4}
+	sort.Sort(byHash(resultVotes))
+
 	if len(votes) != 2 {
 		t.Errorf("fail to get votes: have %d, want 1", len(votes))
 	}
-	if !votes[0].BlockHash.Equal(vote3.BlockHash) {
-		t.Errorf("get wrong vote: have %s, want %s", votes[0], vote3)
+	if !votes[0].BlockHash.Equal(resultVotes[0].BlockHash) {
+		t.Errorf("get wrong vote: have %s, want %s", votes[0], resultVotes[0])
 	}
-	if !votes[1].BlockHash.Equal(vote4.BlockHash) {
-		t.Errorf("get wrong vote: have %s, want %s", votes[1], vote4)
+	if !votes[1].BlockHash.Equal(resultVotes[1].BlockHash) {
+		t.Errorf("get wrong vote: have %s, want %s", votes[1], resultVotes[1])
 	}
 }
 
diff --git a/dex/handler_test.go b/dex/handler_test.go
index dc1ea1d73..c9e99abc9 100644
--- a/dex/handler_test.go
+++ b/dex/handler_test.go
@@ -198,9 +198,10 @@ func testGetBlockHeaders(t *testing.T, protocol int) {
 	// Run each of the tests and verify the results against the chain
 	for i, tt := range tests {
 		// Collect the headers to expect in the response
-		headers := []*types.Header{}
+		headers := []*types.HeaderWithGovState{}
 		for _, hash := range tt.expect {
-			headers = append(headers, pm.blockchain.GetBlockByHash(hash).Header())
+			headers = append(headers, &types.HeaderWithGovState{
+				Header: pm.blockchain.GetBlockByHash(hash).Header()})
 		}
 		// Send the hash request and verify the response
 		p2p.Send(peer.app, 0x03, tt.query)
diff --git a/dex/helper_test.go b/dex/helper_test.go
index f57f3eff4..abe0f6b89 100644
--- a/dex/helper_test.go
+++ b/dex/helper_test.go
@@ -125,7 +125,7 @@ func newTestProtocolManager(mode downloader.SyncMode, blocks int, generator func
 		notarySetFunc: func(uint64, uint32) (map[string]struct{}, error) { return nil, nil },
 	}
 
-	pm, err := NewProtocolManager(gspec.Config, mode, DefaultConfig.NetworkId, evmux, &testTxPool{added: newtx}, engine, blockchain, db, tgov)
+	pm, err := NewProtocolManager(gspec.Config, mode, DefaultConfig.NetworkId, evmux, &testTxPool{added: newtx}, engine, blockchain, db, true, tgov)
 	if err != nil {
 		return nil, nil, err
 	}
@@ -223,6 +223,10 @@ func (g *testGovernance) DKGSet(round uint64) (map[string]struct{}, error) {
 	return g.dkgSetFunc(round)
 }
 
+func (g *testGovernance) GetRoundHeight(round uint64) uint64 {
+	return 0
+}
+
 // testPeer is a simulated peer to allow testing direct network calls.
 type testPeer struct {
 	net p2p.MsgReadWriter // Network layer reader/writer to simulate remote messaging
diff --git a/dex/protocol_test.go b/dex/protocol_test.go
index 64d7fa3e7..c2f9c00b2 100644
--- a/dex/protocol_test.go
+++ b/dex/protocol_test.go
@@ -433,12 +433,14 @@ func TestRecvVote(t *testing.T) {
 	defer p.close()
 
 	vote := coreTypes.Vote{
-		ProposerID: coreTypes.NodeID{coreCommon.Hash{1, 2, 3}},
-		Period:     10,
-		Position: coreTypes.Position{
-			ChainID: 11,
-			Round:   12,
-			Height:  13,
+		VoteHeader: coreTypes.VoteHeader{
+			ProposerID: coreTypes.NodeID{coreCommon.Hash{1, 2, 3}},
+			Period:     10,
+			Position: coreTypes.Position{
+				ChainID: 11,
+				Round:   12,
+				Height:  13,
+			},
 		},
 		Signature: coreCrypto.Signature{
 			Type:      "123",
@@ -468,12 +470,14 @@ func TestSendVote(t *testing.T) {
 	defer pm.Stop()
 
 	vote := coreTypes.Vote{
-		ProposerID: coreTypes.NodeID{coreCommon.Hash{1, 2, 3}},
-		Period:     10,
-		Position: coreTypes.Position{
-			ChainID: 1,
-			Round:   10,
-			Height:  13,
+		VoteHeader: coreTypes.VoteHeader{
+			ProposerID: coreTypes.NodeID{coreCommon.Hash{1, 2, 3}},
+			Period:     10,
+			Position: coreTypes.Position{
+				ChainID: 1,
+				Round:   10,
+				Height:  13,
+			},
 		},
 		Signature: coreCrypto.Signature{
 			Type:      "123",
@@ -663,12 +667,14 @@ func TestRecvAgreement(t *testing.T) {
 
 	// TODO(sonic): polish this
 	vote := coreTypes.Vote{
-		ProposerID: coreTypes.NodeID{coreCommon.Hash{1, 2, 3}},
-		Period:     10,
-		Position: coreTypes.Position{
-			ChainID: 1,
-			Round:   10,
-			Height:  13,
+		VoteHeader: coreTypes.VoteHeader{
+			ProposerID: coreTypes.NodeID{coreCommon.Hash{1, 2, 3}},
+			Period:     10,
+			Position: coreTypes.Position{
+				ChainID: 1,
+				Round:   10,
+				Height:  13,
+			},
 		},
 		Signature: coreCrypto.Signature{
 			Type:      "123",
@@ -706,12 +712,14 @@ func TestSendAgreement(t *testing.T) {
 
 	// TODO(sonic): polish this
 	vote := coreTypes.Vote{
-		ProposerID: coreTypes.NodeID{coreCommon.Hash{1, 2, 3}},
-		Period:     10,
-		Position: coreTypes.Position{
-			ChainID: 1,
-			Round:   10,
-			Height:  13,
+		VoteHeader: coreTypes.VoteHeader{
+			ProposerID: coreTypes.NodeID{coreCommon.Hash{1, 2, 3}},
+			Period:     10,
+			Position: coreTypes.Position{
+				ChainID: 1,
+				Round:   10,
+				Height:  13,
+			},
 		},
 		Signature: coreCrypto.Signature{
 			Type:      "123",