From c2b8cb226eaef2d8b16e2c2048b696eeb448f840 Mon Sep 17 00:00:00 2001
From: Zsolt Felfoldi <zsfelfoldi@gmail.com>
Date: Wed, 17 Jan 2024 17:37:06 +0100
Subject: [PATCH] beacon/light/sync: test Server.Fail calls

---
 beacon/light/sync/test_helpers.go     | 35 +++++++++++++++++++++++++--
 beacon/light/sync/update_sync_test.go |  4 ++-
 2 files changed, 36 insertions(+), 3 deletions(-)

diff --git a/beacon/light/sync/test_helpers.go b/beacon/light/sync/test_helpers.go
index 5d055bd78f..8c0cd1075d 100644
--- a/beacon/light/sync/test_helpers.go
+++ b/beacon/light/sync/test_helpers.go
@@ -24,9 +24,14 @@ import (
 	"github.com/ethereum/go-ethereum/beacon/types"
 )
 
-type TestServer struct{ ID int }
+type TestServer struct {
+	ts *TestScheduler
+	ID int
+}
 
-func (ts *TestServer) Fail(desc string) {}
+func (s *TestServer) Fail(desc string) {
+	s.ts.serverFail(s)
+}
 
 type TestScheduler struct {
 	t         *testing.T
@@ -35,6 +40,8 @@ type TestScheduler struct {
 	servers   []request.Server
 	allowance map[request.Server]int
 	sent      map[int]request.RequestWithID
+	testIndex int
+	expFail   map[request.Server]int // expected Server.Fail calls during next Run
 	lastId    request.ID
 }
 
@@ -43,13 +50,24 @@ func NewTestScheduler(t *testing.T, module request.Module) *TestScheduler {
 		t:         t,
 		module:    module,
 		allowance: make(map[request.Server]int),
+		expFail:   make(map[request.Server]int),
 		sent:      make(map[int]request.RequestWithID),
 	}
 }
 
 func (ts *TestScheduler) Run(testIndex int, expServer request.Server, expReq request.Request) {
+	ts.testIndex = testIndex
 	ts.module.Process(ts.events)
 	ts.events = nil
+
+	for server, count := range ts.expFail {
+		delete(ts.expFail, server)
+		if count == 0 {
+			continue
+		}
+		ts.t.Errorf("Missing %d Server.Fail(s) from server %d in test case #%d", count, server.(*TestServer).ID, testIndex)
+	}
+
 	expReqWithID := request.RequestWithID{
 		ServerAndID: request.ServerAndID{Server: expServer, ID: ts.lastId + 1},
 		Request:     expReq,
@@ -100,6 +118,7 @@ func (ts *TestScheduler) RequestEvent(evType *request.EventType, testIndex int,
 }
 
 func (ts *TestScheduler) AddServer(server request.Server, allowance int) {
+	server.(*TestServer).ts = ts
 	ts.servers = append(ts.servers, server)
 	ts.allowance[server] = allowance
 	ts.ServerEvent(request.EvRegistered, server, nil)
@@ -122,6 +141,18 @@ func (ts *TestScheduler) AddAllowance(server request.Server, allowance int) {
 	ts.allowance[server] += allowance
 }
 
+func (ts *TestScheduler) ExpFail(server request.Server) {
+	ts.expFail[server]++
+}
+
+func (ts *TestScheduler) serverFail(server request.Server) {
+	if ts.expFail[server] == 0 {
+		ts.t.Errorf("Unexpected Server.Fail from server %d in test case #%d", server.(*TestServer).ID, ts.testIndex)
+		return
+	}
+	ts.expFail[server]--
+}
+
 func (ts *TestScheduler) tryRequest(testIndex int, requestFn func(server request.Server) (request.Request, float32)) (request.RequestWithID, bool) {
 	var (
 		bestServer request.Server
diff --git a/beacon/light/sync/update_sync_test.go b/beacon/light/sync/update_sync_test.go
index af966572e8..412654b402 100644
--- a/beacon/light/sync/update_sync_test.go
+++ b/beacon/light/sync/update_sync_test.go
@@ -40,8 +40,9 @@ func TestCheckpointInit(t *testing.T) {
 	ts.RequestEvent(request.EvTimeout, 1, nil)
 	ts.Run(2, testServer2, ReqCheckpointData(checkpointHash))
 
-	// invalid response to server 2; expect init state to still be false
+	// invalid response from server 2; expect init state to still be false
 	ts.RequestEvent(request.EvResponse, 2, &types.BootstrapData{Header: types.Header{Slot: 123456}})
+	ts.ExpFail(testServer2)
 	ts.Run(3, nil, nil)
 	chain.ExpInit(t, false)
 
@@ -172,6 +173,7 @@ func TestUpdateSyncDifferentHeads(t *testing.T) {
 	req1x := ts.Request(1)
 	req1x.Request = ReqUpdates{FirstPeriod: 10, Count: 5}
 	ts.RequestEvent(request.EvResponse, 1, testRespUpdate(req1x))
+	ts.ExpFail(testServer3)
 	ts.Run(5, nil, nil)
 	// expect no progress of chain head
 	chain.ExpNextSyncPeriod(t, 15)