X-Git-Url: http://git.osdn.net/view?p=bytom%2Fvapor.git;a=blobdiff_plain;f=netsync%2Fconsensusmgr%2Fblock_fetcher_test.go;h=96703516c1fb9449facc13088b06bf0ef8f2df6b;hp=8d0df223fb942864c252b297b9963efc9a194a0b;hb=refs%2Fheads%2Fblock_fetcher;hpb=a5ca0c085e1389ce83e2f877c3a72a978ac42ab3

diff --git a/netsync/consensusmgr/block_fetcher_test.go b/netsync/consensusmgr/block_fetcher_test.go
index 8d0df223..96703516 100644
--- a/netsync/consensusmgr/block_fetcher_test.go
+++ b/netsync/consensusmgr/block_fetcher_test.go
@@ -127,7 +127,7 @@ func TestBlockFetcher(t *testing.T) {
 		},
 	}
 	fetcher := newBlockFetcher(newChain(), peers)
-
+	go fetcher.blockProcessorLoop()
 	for i, c := range testCase {
 		fetcher.processNewBlock(c.blockMsg)
 		time.Sleep(10 * time.Millisecond)
@@ -137,3 +137,159 @@ func TestBlockFetcher(t *testing.T) {
 		}
 	}
 }
+
+func TestAddBlockMsg(t *testing.T) {
+	peers := peers.NewPeerSet(&peerMgr{})
+	testPeer := "peer1"
+	testCase := []struct {
+		blocksMsg  []*blockMsg
+		limit      int
+		queueSize  int
+		msgSetSize int
+		msgCounter int
+	}{
+		//normal test
+		{
+			blocksMsg: []*blockMsg{
+				{
+					block: &types.Block{
+						BlockHeader: types.BlockHeader{
+							Height: 100,
+						},
+					},
+					peerID: testPeer,
+				},
+				{
+					block: &types.Block{
+						BlockHeader: types.BlockHeader{
+							Height: 101,
+						},
+					},
+					peerID: testPeer,
+				},
+				{
+					block: &types.Block{
+						BlockHeader: types.BlockHeader{
+							Height: 102,
+						},
+					},
+					peerID: testPeer,
+				},
+			},
+			limit:      5,
+			queueSize:  3,
+			msgSetSize: 3,
+			msgCounter: 3,
+		},
+		// test DOS
+		{
+			blocksMsg: []*blockMsg{
+				{
+					block: &types.Block{
+						BlockHeader: types.BlockHeader{
+							Version: 1,
+							Height:  100,
+						},
+					},
+					peerID: testPeer,
+				},
+				{
+					block: &types.Block{
+						BlockHeader: types.BlockHeader{
+							Version: 2,
+							Height:  100,
+						},
+					},
+					peerID: testPeer,
+				},
+				{
+					block: &types.Block{
+						BlockHeader: types.BlockHeader{
+							Version: 3,
+							Height:  100,
+						},
+					},
+					peerID: testPeer,
+				},
+				{
+					block: &types.Block{
+						BlockHeader: types.BlockHeader{
+							Version: 4,
+							Height:  100,
+						},
+					},
+					peerID: testPeer,
+				},
+			},
+			limit:      3,
+			queueSize:  3,
+			msgSetSize: 3,
+			msgCounter: 3,
+		},
+
+		// test msg height does not meet the requirements
+		{
+			blocksMsg: []*blockMsg{
+				{
+					block: &types.Block{
+						BlockHeader: types.BlockHeader{
+							Version: 1,
+							Height:  98,
+						},
+					},
+					peerID: testPeer,
+				},
+				{
+					block: &types.Block{
+						BlockHeader: types.BlockHeader{
+							Version: 2,
+							Height:  97,
+						},
+					},
+					peerID: testPeer,
+				},
+				{
+					block: &types.Block{
+						BlockHeader: types.BlockHeader{
+							Version: 3,
+							Height:  164,
+						},
+					},
+					peerID: testPeer,
+				},
+				{
+					block: &types.Block{
+						BlockHeader: types.BlockHeader{
+							Version: 4,
+							Height:  165,
+						},
+					},
+					peerID: testPeer,
+				},
+			},
+			limit:      5,
+			queueSize:  0,
+			msgSetSize: 0,
+			msgCounter: 0,
+		},
+	}
+
+	for i, c := range testCase {
+		fetcher := newBlockFetcher(newChain(), peers)
+		for _, msg := range c.blocksMsg {
+			fetcher.add(msg, c.limit)
+		}
+
+		if fetcher.queue.Size() != c.queueSize {
+			t.Fatalf("index: %d queue size err got %d: want %d", i, fetcher.queue.Size(), c.queueSize)
+		}
+
+		if len(fetcher.msgSet) != c.msgSetSize {
+			t.Fatalf("index: %d msg set size err got %d: want %d", i, len(fetcher.msgSet), c.msgSetSize)
+		}
+
+		if fetcher.msgCounter[testPeer] != c.msgCounter {
+			t.Fatalf("index: %d peer msg counter err got %d: want %d", i, fetcher.msgCounter[testPeer], c.msgCounter)
+		}
+	}
+}