From 0bba2226fe49c9d41f31029c5bb699446a5cd9aa Mon Sep 17 00:00:00 2001
From: Andrew <i.am.andrew.cullen@gmail.com>
Date: Wed, 24 Apr 2024 16:05:51 +0100
Subject: [PATCH] fix underflows in rewards tests waiting

---
 tools/docker-network/tests/rewards_test.go | 38 +++++++++++++---------
 tools/docker-network/tests/utils.go        | 10 +++---
 2 files changed, 28 insertions(+), 20 deletions(-)

diff --git a/tools/docker-network/tests/rewards_test.go b/tools/docker-network/tests/rewards_test.go
index 10a4303a1..35ab84c60 100644
--- a/tools/docker-network/tests/rewards_test.go
+++ b/tools/docker-network/tests/rewards_test.go
@@ -68,16 +68,17 @@ func Test_ValidatorRewards(t *testing.T) {
 	d.AssertCommittee(currentEpoch+2, append(d.AccountsFromNodes(d.Nodes("V1", "V3", "V2", "V4")...), goodAccountAddrBech32, lazyAccountAddrBech32))
 
 	// issue validation blocks to have performance
-	currentSlot := clt.CommittedAPI().TimeProvider().CurrentSlot()
-	slotToWait := claimingSlot - currentSlot
-	secToWait := time.Duration(slotToWait) * time.Duration(slotsDuration) * time.Second
-	fmt.Println("Wait for ", secToWait, "until expected slot: ", claimingSlot)
+	if currentSlot := clt.CommittedAPI().TimeProvider().CurrentSlot(); currentSlot < claimingSlot {
+		slotToWait := claimingSlot - currentSlot
+		secToWait := time.Duration(slotToWait) * time.Duration(slotsDuration) * time.Second
+		fmt.Println("Wait for ", secToWait, "until expected slot: ", claimingSlot)
 
-	var wg sync.WaitGroup
-	issueValidationBlockInBackground(&wg, goodWallet, currentSlot, claimingSlot, 5)
-	issueValidationBlockInBackground(&wg, lazyWallet, currentSlot, claimingSlot, 1)
+		var wg sync.WaitGroup
+		issueValidationBlockInBackground(&wg, goodWallet, currentSlot, claimingSlot, 5)
+		issueValidationBlockInBackground(&wg, lazyWallet, currentSlot, claimingSlot, 1)
 
-	wg.Wait()
+		wg.Wait()
+	}
 
 	// claim rewards that put to the account output
 	d.AwaitCommitment(claimingSlot)
@@ -135,10 +136,12 @@ func Test_DelegatorRewards(t *testing.T) {
 	// wait until next epoch so the rewards can be claimed
 	//nolint:forcetypeassert
 	expectedSlot := clt.CommittedAPI().TimeProvider().EpochStart(delegationOutputData.Output.(*iotago.DelegationOutput).StartEpoch + 2)
-	slotToWait := expectedSlot - clt.CommittedAPI().TimeProvider().CurrentSlot()
-	secToWait := time.Duration(slotToWait) * time.Duration(clt.CommittedAPI().ProtocolParameters().SlotDurationInSeconds()) * time.Second
-	fmt.Println("Wait for ", secToWait, "until expected slot: ", expectedSlot)
-	time.Sleep(secToWait)
+	if currentSlot := clt.CommittedAPI().TimeProvider().CurrentSlot(); currentSlot < expectedSlot {
+		slotToWait := expectedSlot - currentSlot
+		secToWait := time.Duration(slotToWait) * time.Duration(clt.CommittedAPI().ProtocolParameters().SlotDurationInSeconds()) * time.Second
+		fmt.Println("Wait for ", secToWait, "until expected slot: ", expectedSlot)
+		time.Sleep(secToWait)
+	}
 
 	// claim rewards that put to an basic output
 	rewardsOutputID := d.ClaimRewardsForDelegator(ctx, delegatorWallet, delegationOutputData)
@@ -204,10 +207,13 @@ func Test_DelayedClaimingRewards(t *testing.T) {
 
 		// wait until next epoch to destroy the delegation
 		expectedSlot := clt.CommittedAPI().TimeProvider().EpochStart(delegationEndEpoch)
-		slotToWait := expectedSlot - clt.CommittedAPI().TimeProvider().CurrentSlot()
-		secToWait := time.Duration(slotToWait) * time.Duration(clt.CommittedAPI().ProtocolParameters().SlotDurationInSeconds()) * time.Second
-		fmt.Println("Wait for ", secToWait, "until expected slot: ", expectedSlot)
-		time.Sleep(secToWait)
+		if currentSlot := delegatorWallet.CurrentSlot(); currentSlot < expectedSlot {
+			slotToWait := expectedSlot - currentSlot
+			secToWait := time.Duration(slotToWait) * time.Duration(clt.CommittedAPI().ProtocolParameters().SlotDurationInSeconds()) * time.Second
+			fmt.Println("Wait for ", secToWait, "until expected slot: ", expectedSlot)
+			time.Sleep(secToWait)
+		}
+		fmt.Println("Claim rewards for delegator")
 		d.ClaimRewardsForDelegator(ctx, delegatorWallet, delegationOutputData)
 	}
 
diff --git a/tools/docker-network/tests/utils.go b/tools/docker-network/tests/utils.go
index 88568af2e..2e81ac8c8 100644
--- a/tools/docker-network/tests/utils.go
+++ b/tools/docker-network/tests/utils.go
@@ -119,10 +119,12 @@ func (d *DockerTestFramework) AssertCommittee(expectedEpoch iotago.EpochIndex, e
 	expectedSlotStart := testAPI.TimeProvider().EpochStart(expectedEpoch)
 	require.Greater(d.Testing, expectedSlotStart, status.LatestAcceptedBlockSlot)
 
-	slotToWait := expectedSlotStart - status.LatestAcceptedBlockSlot
-	secToWait := time.Duration(slotToWait) * time.Duration(testAPI.ProtocolParameters().SlotDurationInSeconds()) * time.Second
-	fmt.Println("Wait for ", secToWait, "until expected epoch: ", expectedEpoch)
-	time.Sleep(secToWait)
+	if status.LatestAcceptedBlockSlot < expectedSlotStart {
+		slotToWait := expectedSlotStart - status.LatestAcceptedBlockSlot
+		secToWait := time.Duration(slotToWait) * time.Duration(testAPI.ProtocolParameters().SlotDurationInSeconds()) * time.Second
+		fmt.Println("Wait for ", secToWait, "until expected epoch: ", expectedEpoch)
+		time.Sleep(secToWait)
+	}
 
 	d.Eventually(func() error {
 		for _, node := range d.Nodes() {