From 1d031148ac4ead67e77ca6140156aa858efa16c4 Mon Sep 17 00:00:00 2001
From: jkrvivian <jkrvivian@gmail.com>
Date: Wed, 15 May 2024 15:51:35 +0800
Subject: [PATCH] Await Accepted slot in issueCandidacy and
 issueValidationBlock in Test_ValidatorRewards

---
 .../tests/dockertestframework/awaits.go       | 26 +++++++++++++++++++
 tools/docker-network/tests/rewards_test.go    | 17 ++----------
 2 files changed, 28 insertions(+), 15 deletions(-)

diff --git a/tools/docker-network/tests/dockertestframework/awaits.go b/tools/docker-network/tests/dockertestframework/awaits.go
index 26685aa15..747d9940b 100644
--- a/tools/docker-network/tests/dockertestframework/awaits.go
+++ b/tools/docker-network/tests/dockertestframework/awaits.go
@@ -69,6 +69,32 @@ func (d *DockerTestFramework) AwaitTransactionFailure(ctx context.Context, txID
 	})
 }
 
+func (d *DockerTestFramework) AwaitSlot(targetSlot iotago.SlotIndex) {
+	currentSlot := d.NodeStatus("V1").LatestAcceptedBlockSlot
+
+	if targetSlot <= currentSlot {
+		return
+	}
+
+	// we wait at max "targetSlot - currentSlot" times * slot duration
+	deadline := time.Duration(d.defaultWallet.Client.CommittedAPI().ProtocolParameters().SlotDurationInSeconds()) * time.Second
+	if currentSlot < targetSlot {
+		deadline *= time.Duration(targetSlot - currentSlot)
+	}
+
+	// give some extra time for peering etc
+	deadline += 30 * time.Second
+
+	d.EventuallyWithDurations(func() error {
+		latestAcceptedSlot := d.NodeStatus("V1").LatestAcceptedBlockSlot
+		if targetSlot > latestAcceptedSlot {
+			return ierrors.Errorf("current slot %d is not reached yet, current slot %d", targetSlot, latestAcceptedSlot)
+		}
+
+		return nil
+	}, deadline, 1*time.Second)
+}
+
 func (d *DockerTestFramework) AwaitCommitment(targetSlot iotago.SlotIndex) {
 	currentCommittedSlot := d.NodeStatus("V1").LatestCommitmentID.Slot()
 
diff --git a/tools/docker-network/tests/rewards_test.go b/tools/docker-network/tests/rewards_test.go
index 74072204a..d3ba71372 100644
--- a/tools/docker-network/tests/rewards_test.go
+++ b/tools/docker-network/tests/rewards_test.go
@@ -298,7 +298,7 @@ func issueCandidacyAnnouncementsInBackground(ctx context.Context, d *dockertestf
 			}
 
 			// wait until the epoch start is reached
-			d.AwaitCommitment(d.DefaultWallet().Client.CommittedAPI().TimeProvider().EpochStart(epoch))
+			d.AwaitSlot(d.DefaultWallet().Client.CommittedAPI().TimeProvider().EpochStart(epoch))
 			if ctx.Err() != nil {
 				// context is canceled
 				return
@@ -307,12 +307,6 @@ func issueCandidacyAnnouncementsInBackground(ctx context.Context, d *dockertestf
 			fmt.Println("Issuing candidacy payload for account", wallet.BlockIssuer.AccountData.ID, "in epoch", epoch, "...")
 			committedAPI := d.DefaultWallet().Client.CommittedAPI()
 
-			// check if we are still in the epoch
-			currentCommittedSlot := d.NodeStatus("V1").LatestCommitmentID.Slot()
-			currentEpoch := committedAPI.TimeProvider().EpochFromSlot(currentCommittedSlot)
-
-			require.Equal(d.Testing, epoch, currentEpoch, "epoch mismatch")
-
 			// the candidacy announcement needs to be done before the nearing threshold
 			epochEndSlot := committedAPI.TimeProvider().EpochEnd(epoch)
 			maxRegistrationSlot := epochEndSlot - committedAPI.ProtocolParameters().EpochNearingThreshold()
@@ -342,19 +336,12 @@ func issueValidationBlocksInBackground(ctx context.Context, d *dockertestframewo
 			}
 
 			// wait until the slot is reached
-			d.AwaitCommitment(slot)
+			d.AwaitSlot(slot)
 			if ctx.Err() != nil {
 				// context is canceled
 				return
 			}
 
-			// check if we are still in the slot
-			currentCommittedSlot := d.NodeStatus("V1").LatestCommitmentID.Slot()
-			if slot < currentCommittedSlot {
-				// slot is already committed, no need to issue validation blocks
-				continue
-			}
-
 			ts := time.Now()
 			for validationBlockNr := range blocksPerSlot {
 				if ctx.Err() != nil {