@@ -68,6 +68,14 @@ func (i poolTestPeer) freeClientId() string {
6868
6969func (i poolTestPeer ) updateCapacity (uint64 ) {}
7070
71+ type poolTestPeerWithCap struct {
72+ poolTestPeer
73+
74+ cap uint64
75+ }
76+
77+ func (i * poolTestPeerWithCap ) updateCapacity (cap uint64 ) { i .cap = cap }
78+
7179func testClientPool (t * testing.T , connLimit , clientCount , paidCount int , randomDisconnect bool ) {
7280 rand .Seed (time .Now ().UnixNano ())
7381 var (
@@ -308,9 +316,9 @@ func TestFreeClientKickedOut(t *testing.T) {
308316
309317 for i := 0 ; i < 10 ; i ++ {
310318 pool .connect (poolTestPeer (i ), 1 )
311- clock .Run (100 * time .Millisecond )
319+ clock .Run (time .Millisecond )
312320 }
313- if pool .connect (poolTestPeer (11 ), 1 ) {
321+ if pool .connect (poolTestPeer (10 ), 1 ) {
314322 t .Fatalf ("New free client should be rejected" )
315323 }
316324 clock .Run (5 * time .Minute )
@@ -320,8 +328,8 @@ func TestFreeClientKickedOut(t *testing.T) {
320328 for i := 0 ; i < 10 ; i ++ {
321329 select {
322330 case id := <- kicked :
323- if id != i {
324- t .Fatalf ("Kicked client mismatch, want %v, got %v" , i , id )
331+ if id >= 10 {
332+ t .Fatalf ("Old client should be kicked, now got: %d" , id )
325333 }
326334 case <- time .NewTimer (time .Second ).C :
327335 t .Fatalf ("timeout" )
@@ -364,11 +372,20 @@ func TestDowngradePriorityClient(t *testing.T) {
364372 pool .setLimits (10 , uint64 (10 )) // Total capacity limit is 10
365373 pool .setPriceFactors (priceFactors {1 , 0 , 1 }, priceFactors {1 , 0 , 1 })
366374
367- pool .addBalance (poolTestPeer (0 ).ID (), uint64 (time .Minute ), false )
368- pool .connect (poolTestPeer (0 ), 10 )
375+ p := & poolTestPeerWithCap {
376+ poolTestPeer : poolTestPeer (0 ),
377+ }
378+ pool .addBalance (p .ID (), uint64 (time .Minute ), false )
379+ pool .connect (p , 10 )
380+ if p .cap != 10 {
381+ t .Fatalf ("The capcacity of priority peer hasn't been updated, got: %d" , p .cap )
382+ }
383+
369384 clock .Run (time .Minute ) // All positive balance should be used up.
370385 time .Sleep (300 * time .Millisecond ) // Ensure the callback is called
371-
386+ if p .cap != 1 {
387+ t .Fatalf ("The capcacity of peer should be downgraded, got: %d" , p .cap )
388+ }
372389 pb := pool .ndb .getOrNewPB (poolTestPeer (0 ).ID ())
373390 if pb .value != 0 {
374391 t .Fatalf ("Positive balance mismatch, want %v, got %v" , 0 , pb .value )
0 commit comments