les: update unit tests

les/clientpool.go

@@ -459,6 +459,9 @@ func (f *clientPool) addBalance(id enode.ID, amount uint64, setTotal bool) {
 		defer func() {
 			c.balanceTracker.setBalance(pb.value, negBalance)
 			if !c.priority && pb.value > 0 {
+				// The capacity should be adjusted based on the requirement,
+				// but we have no idea about the new capacity, need a second
+				// call to udpate it.
 				c.priority = true
 				c.balanceTracker.addCallback(balanceCallbackZero, 0, func() { f.balanceExhausted(id) })
 			}

les/clientpool_test.go

@@ -68,6 +68,14 @@ func (i poolTestPeer) freeClientId() string {
 
 func (i poolTestPeer) updateCapacity(uint64) {}
 
+type poolTestPeerWithCap struct {
+	poolTestPeer
+
+	cap uint64
+}
+
+func (i *poolTestPeerWithCap) updateCapacity(cap uint64) { i.cap = cap }
+
 func testClientPool(t *testing.T, connLimit, clientCount, paidCount int, randomDisconnect bool) {
 	rand.Seed(time.Now().UnixNano())
 	var (
@@ -364,11 +372,20 @@ func TestDowngradePriorityClient(t *testing.T) {
 	pool.setLimits(10, uint64(10)) // Total capacity limit is 10
 	pool.setPriceFactors(priceFactors{1, 0, 1}, priceFactors{1, 0, 1})
 
-	pool.addBalance(poolTestPeer(0).ID(), uint64(time.Minute), false)
+	p := &poolTestPeerWithCap{
-	pool.connect(poolTestPeer(0), 10)
+		poolTestPeer: poolTestPeer(0),
+	}
+	pool.addBalance(p.ID(), uint64(time.Minute), false)
+	pool.connect(p, 10)
+	if p.cap != 10 {
+		t.Fatalf("The capcacity of priority peer hasn't been updated, got: %d", p.cap)
+	}
+
 	clock.Run(time.Minute)             // All positive balance should be used up.
 	time.Sleep(300 * time.Millisecond) // Ensure the callback is called
-
+	if p.cap != 1 {
+		t.Fatalf("The capcacity of peer should be downgraded, got: %d", p.cap)
+	}
 	pb := pool.ndb.getOrNewPB(poolTestPeer(0).ID())
 	if pb.value != 0 {
 		t.Fatalf("Positive balance mismatch, want %v, got %v", 0, pb.value)

les/odr_test.go

@@ -202,7 +202,7 @@ func testOdr(t *testing.T, protocol int, expFail uint64, checkCached bool, fn od
 			bhash := rawdb.ReadCanonicalHash(server.db, i)
 			b1 := fn(light.NoOdr, server.db, server.handler.server.chainConfig, server.handler.blockchain, nil, bhash)
 
-			ctx, cancel := context.WithTimeout(context.Background(), 200*time.Millisecond)
+			ctx, cancel := context.WithTimeout(context.Background(), time.Second)
 			b2 := fn(ctx, client.db, client.handler.backend.chainConfig, nil, client.handler.backend.blockchain, bhash)
 			cancel()