les: improved header fetcher and server statistics

1 files changed, 5 insertions, 3 deletions

diff --git a/les/odr_test.go b/les/odr_test.go
index 27e3b283d..685435996 100644
--- a/les/odr_test.go
+++ b/les/odr_test.go
@@ -160,6 +160,8 @@ func testOdr(t *testing.T, protocol int, expFail uint64, fn odrTestFn) {
 	pm, db, odr := newTestProtocolManagerMust(t, false, 4, testChainGen)
 	lpm, ldb, odr := newTestProtocolManagerMust(t, true, 0, nil)
 	_, err1, lpeer, err2 := newTestPeerPair("peer", protocol, pm, lpm)
+	pool := (*testServerPool)(lpeer)
+	odr.serverPool = pool
 	select {
 	case <-time.After(time.Millisecond * 100):
 	case err := <-err1:
@@ -188,13 +190,13 @@ func testOdr(t *testing.T, protocol int, expFail uint64, fn odrTestFn) {
 	}
 
 	// temporarily remove peer to test odr fails
-	odr.UnregisterPeer(lpeer)
+	odr.serverPool = nil
 	// expect retrievals to fail (except genesis block) without a les peer
 	test(expFail)
-	odr.RegisterPeer(lpeer)
+	odr.serverPool = pool
 	// expect all retrievals to pass
 	test(5)
-	odr.UnregisterPeer(lpeer)
+	odr.serverPool = nil
 	// still expect all retrievals to pass, now data should be cached locally
 	test(5)
 }