1 files changed, 60 insertions, 31 deletions

diff --git a/eth/downloader/downloader_test.go b/eth/downloader/downloader_test.go
index 5518163ca..872ea02eb 100644
--- a/eth/downloader/downloader_test.go
+++ b/eth/downloader/downloader_test.go
@@ -8,8 +8,6 @@ import (
 
 	"github.com/ethereum/go-ethereum/common"
 	"github.com/ethereum/go-ethereum/core/types"
-	"github.com/ethereum/go-ethereum/logger"
-	"github.com/ethereum/go-ethereum/logger/glog"
 )
 
 var knownHash = common.Hash{1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0}
@@ -25,36 +23,47 @@ func createHashes(start, amount int) (hashes []common.Hash) {
 	return
 }
 
+func createBlock(i int, prevHash, hash common.Hash) *types.Block {
+	header := &types.Header{Number: big.NewInt(int64(i))}
+	block := types.NewBlockWithHeader(header)
+	block.HeaderHash = hash
+	block.ParentHeaderHash = knownHash
+	return block
+}
+
 func createBlocksFromHashes(hashes []common.Hash) map[common.Hash]*types.Block {
 	blocks := make(map[common.Hash]*types.Block)
+
 	for i, hash := range hashes {
-		header := &types.Header{Number: big.NewInt(int64(i))}
-		blocks[hash] = types.NewBlockWithHeader(header)
-		blocks[hash].HeaderHash = hash
+		blocks[hash] = createBlock(len(hashes)-i, knownHash, hash)
 	}
 
 	return blocks
 }
 
 type downloadTester struct {
-	downloader *Downloader
-	hashes     []common.Hash
-	blocks     map[common.Hash]*types.Block
-	t          *testing.T
-	pcount     int
-	done       chan bool
-
-	insertedBlocks int
+	downloader   *Downloader
+	hashes       []common.Hash
+	blocks       map[common.Hash]*types.Block
+	t            *testing.T
+	pcount       int
+	done         chan bool
+	activePeerId string
 }
 
 func newTester(t *testing.T, hashes []common.Hash, blocks map[common.Hash]*types.Block) *downloadTester {
 	tester := &downloadTester{t: t, hashes: hashes, blocks: blocks, done: make(chan bool)}
-	downloader := New(tester.hasBlock, tester.insertChain)
+	downloader := New(tester.hasBlock, tester.getBlock)
 	tester.downloader = downloader
 
 	return tester
 }
 
+func (dl *downloadTester) sync(peerId string, hash common.Hash) error {
+	dl.activePeerId = peerId
+	return dl.downloader.Synchronise(peerId, hash)
+}
+
 func (dl *downloadTester) hasBlock(hash common.Hash) bool {
 	if knownHash == hash {
 		return true
@@ -62,14 +71,12 @@ func (dl *downloadTester) hasBlock(hash common.Hash) bool {
 	return false
 }
 
-func (dl *downloadTester) insertChain(blocks types.Blocks) (int, error) {
-	dl.insertedBlocks += len(blocks)
-
-	return 0, nil
+func (dl *downloadTester) getBlock(hash common.Hash) *types.Block {
+	return dl.blocks[knownHash]
 }
 
 func (dl *downloadTester) getHashes(hash common.Hash) error {
-	dl.downloader.hashCh <- dl.hashes
+	dl.downloader.AddHashes(dl.activePeerId, dl.hashes)
 	return nil
 }
 
@@ -102,9 +109,6 @@ func (dl *downloadTester) badBlocksPeer(id string, td *big.Int, hash common.Hash
 }
 
 func TestDownload(t *testing.T) {
-	glog.SetV(logger.Detail)
-	glog.SetToStderr(true)
-
 	minDesiredPeerCount = 4
 	blockTtl = 1 * time.Second
 
@@ -117,21 +121,20 @@ func TestDownload(t *testing.T) {
 	tester.newPeer("peer2", big.NewInt(0), common.Hash{})
 	tester.badBlocksPeer("peer3", big.NewInt(0), common.Hash{})
 	tester.badBlocksPeer("peer4", big.NewInt(0), common.Hash{})
+	tester.activePeerId = "peer1"
 
-	err := tester.downloader.Synchronise("peer1", hashes[0])
+	err := tester.sync("peer1", hashes[0])
 	if err != nil {
 		t.Error("download error", err)
 	}
 
-	if tester.insertedBlocks != targetBlocks {
-		t.Error("expected", targetBlocks, "have", tester.insertedBlocks)
+	inqueue := len(tester.downloader.queue.blocks)
+	if inqueue != targetBlocks {
+		t.Error("expected", targetBlocks, "have", inqueue)
 	}
 }
 
 func TestMissing(t *testing.T) {
-	glog.SetV(logger.Detail)
-	glog.SetToStderr(true)
-
 	targetBlocks := 1000
 	hashes := createHashes(0, 1000)
 	extraHashes := createHashes(1001, 1003)
@@ -143,12 +146,38 @@ func TestMissing(t *testing.T) {
 	hashes = append(extraHashes, hashes[:len(hashes)-1]...)
 	tester.newPeer("peer2", big.NewInt(0), common.Hash{})
 
-	err := tester.downloader.Synchronise("peer1", hashes[0])
+	err := tester.sync("peer1", hashes[0])
+	if err != nil {
+		t.Error("download error", err)
+	}
+
+	inqueue := len(tester.downloader.queue.blocks)
+	if inqueue != targetBlocks {
+		t.Error("expected", targetBlocks, "have", inqueue)
+	}
+}
+
+func TestTaking(t *testing.T) {
+	minDesiredPeerCount = 4
+	blockTtl = 1 * time.Second
+
+	targetBlocks := 1000
+	hashes := createHashes(0, targetBlocks)
+	blocks := createBlocksFromHashes(hashes)
+	tester := newTester(t, hashes, blocks)
+
+	tester.newPeer("peer1", big.NewInt(10000), hashes[0])
+	tester.newPeer("peer2", big.NewInt(0), common.Hash{})
+	tester.badBlocksPeer("peer3", big.NewInt(0), common.Hash{})
+	tester.badBlocksPeer("peer4", big.NewInt(0), common.Hash{})
+
+	err := tester.sync("peer1", hashes[0])
 	if err != nil {
 		t.Error("download error", err)
 	}
 
-	if tester.insertedBlocks != targetBlocks {
-		t.Error("expected", targetBlocks, "have", tester.insertedBlocks)
+	bs1 := tester.downloader.TakeBlocks()
+	if len(bs1) != 1000 {
+		t.Error("expected to take 1000, got", len(bs1))
 	}
 }