go-pulse/les/vflux/client/fillset_test.go
2021-02-19 14:44:16 +01:00

114 lines
2.7 KiB
Go

// Copyright 2020 The go-ethereum Authors
// This file is part of the go-ethereum library.
//
// The go-ethereum library is free software: you can redistribute it and/or modify
// it under the terms of the GNU Lesser General Public License as published by
// the Free Software Foundation, either version 3 of the License, or
// (at your option) any later version.
//
// The go-ethereum library is distributed in the hope that it will be useful,
// but WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
// GNU Lesser General Public License for more details.
//
// You should have received a copy of the GNU Lesser General Public License
// along with the go-ethereum library. If not, see <http://www.gnu.org/licenses/>.
package client
import (
"math/rand"
"testing"
"time"
"github.com/ethereum/go-ethereum/common/mclock"
"github.com/ethereum/go-ethereum/p2p/enode"
"github.com/ethereum/go-ethereum/p2p/enr"
"github.com/ethereum/go-ethereum/p2p/nodestate"
)
type testIter struct {
waitCh chan struct{}
nodeCh chan *enode.Node
node *enode.Node
}
func (i *testIter) Next() bool {
i.waitCh <- struct{}{}
i.node = <-i.nodeCh
return i.node != nil
}
func (i *testIter) Node() *enode.Node {
return i.node
}
func (i *testIter) Close() {}
func (i *testIter) push() {
var id enode.ID
rand.Read(id[:])
i.nodeCh <- enode.SignNull(new(enr.Record), id)
}
func (i *testIter) waiting(timeout time.Duration) bool {
select {
case <-i.waitCh:
return true
case <-time.After(timeout):
return false
}
}
func TestFillSet(t *testing.T) {
ns := nodestate.NewNodeStateMachine(nil, nil, &mclock.Simulated{}, testSetup)
iter := &testIter{
waitCh: make(chan struct{}),
nodeCh: make(chan *enode.Node),
}
fs := NewFillSet(ns, iter, sfTest1)
ns.Start()
expWaiting := func(i int, push bool) {
for ; i > 0; i-- {
if !iter.waiting(time.Second * 10) {
t.Fatalf("FillSet not waiting for new nodes")
}
if push {
iter.push()
}
}
}
expNotWaiting := func() {
if iter.waiting(time.Millisecond * 100) {
t.Fatalf("FillSet unexpectedly waiting for new nodes")
}
}
expNotWaiting()
fs.SetTarget(3)
expWaiting(3, true)
expNotWaiting()
fs.SetTarget(100)
expWaiting(2, true)
expWaiting(1, false)
// lower the target before the previous one has been filled up
fs.SetTarget(0)
iter.push()
expNotWaiting()
fs.SetTarget(10)
expWaiting(4, true)
expNotWaiting()
// remove all previosly set flags
ns.ForEach(sfTest1, nodestate.Flags{}, func(node *enode.Node, state nodestate.Flags) {
ns.SetState(node, nodestate.Flags{}, sfTest1, 0)
})
// now expect FillSet to fill the set up again with 10 new nodes
expWaiting(10, true)
expNotWaiting()
fs.Close()
ns.Stop()
}