From bf4393c786ebacf5092bb7584a353cfe4a486438 Mon Sep 17 00:00:00 2001
From: jkrvivian <jkrvivian@gmail.com>
Date: Thu, 9 May 2024 14:26:16 +0800
Subject: [PATCH] Fix waiting time in `issueCandidacyPayloadInBackground`

---
 tools/docker-network/tests/api_core_test.go   |  2 +-
 .../tests/dockertestframework/validator.go    |  4 +--
 tools/docker-network/tests/rewards_test.go    | 27 +++++++++++--------
 3 files changed, 19 insertions(+), 14 deletions(-)

diff --git a/tools/docker-network/tests/api_core_test.go b/tools/docker-network/tests/api_core_test.go
index 5ddbe18f8..5c064ce61 100644
--- a/tools/docker-network/tests/api_core_test.go
+++ b/tools/docker-network/tests/api_core_test.go
@@ -311,7 +311,7 @@ func Test_ValidatorsAPI(t *testing.T) {
 
 			// issue candidacy payload in the next epoch (currentEpoch + 1), in order to issue it before epochNearingThreshold
 			d.AwaitCommitment(clt.CommittedAPI().TimeProvider().EpochEnd(currentEpoch))
-			blkID := d.IssueCandidacyPayloadFromAccount(wallet)
+			blkID := d.IssueCandidacyPayloadFromAccount(ctx, wallet)
 			fmt.Println("Candidacy payload:", blkID.ToHex(), blkID.Slot())
 			d.AwaitCommitment(blkID.Slot())
 		}()
diff --git a/tools/docker-network/tests/dockertestframework/validator.go b/tools/docker-network/tests/dockertestframework/validator.go
index 68c57335f..534c7e21e 100644
--- a/tools/docker-network/tests/dockertestframework/validator.go
+++ b/tools/docker-network/tests/dockertestframework/validator.go
@@ -37,8 +37,8 @@ func (d *DockerTestFramework) StopIssueCandidacyPayload(nodes ...*Node) {
 	require.NoError(d.Testing, err)
 }
 
-func (d *DockerTestFramework) IssueCandidacyPayloadFromAccount(wallet *mock.Wallet) iotago.BlockID {
-	block, err := wallet.CreateAndSubmitBasicBlock(context.TODO(), "candidacy_payload", mock.WithPayload(&iotago.CandidacyAnnouncement{}))
+func (d *DockerTestFramework) IssueCandidacyPayloadFromAccount(ctx context.Context, wallet *mock.Wallet) iotago.BlockID {
+	block, err := wallet.CreateAndSubmitBasicBlock(ctx, "candidacy_payload", mock.WithPayload(&iotago.CandidacyAnnouncement{}))
 	require.NoError(d.Testing, err)
 
 	return block.ID()
diff --git a/tools/docker-network/tests/rewards_test.go b/tools/docker-network/tests/rewards_test.go
index 539e6e7c9..3a3022c25 100644
--- a/tools/docker-network/tests/rewards_test.go
+++ b/tools/docker-network/tests/rewards_test.go
@@ -73,8 +73,7 @@ func Test_ValidatorRewards(t *testing.T) {
 		d,
 		goodWallet,
 		clt.CommittedAPI().TimeProvider().CurrentSlot(),
-		goodClaimingSlot,
-		slotsDuration)
+		goodClaimingSlot)
 
 	// create lazy account
 	lazyWallet, lazyAccountOutputData := d.CreateImplicitAccount(ctx)
@@ -98,8 +97,7 @@ func Test_ValidatorRewards(t *testing.T) {
 		d,
 		lazyWallet,
 		clt.CommittedAPI().TimeProvider().CurrentSlot(),
-		lazyClaimingSlot,
-		slotsDuration)
+		lazyClaimingSlot)
 
 	// make sure the account is in the committee, so it can issue validation blocks
 	goodAccountAddrBech32 := goodAccountData.Address.Bech32(clt.CommittedAPI().ProtocolParameters().Bech32HRP())
@@ -273,19 +271,26 @@ func Test_DelayedClaimingRewards(t *testing.T) {
 	}
 }
 
-func issueCandidacyPayloadInBackground(ctx context.Context, d *dockertestframework.DockerTestFramework, wallet *mock.Wallet, startSlot, endSlot iotago.SlotIndex, slotDuration uint8) {
+func issueCandidacyPayloadInBackground(ctx context.Context, d *dockertestframework.DockerTestFramework, wallet *mock.Wallet, startSlot, endSlot iotago.SlotIndex) {
 	go func() {
 		fmt.Println("Issuing candidacy payloads for account", wallet.BlockIssuer.AccountData.ID, "in the background...")
 		defer fmt.Println("Issuing candidacy payloads for account", wallet.BlockIssuer.AccountData.ID, "in the background......done")
 
 		for i := startSlot; i < endSlot; i++ {
-			if ctx.Err() != nil {
-				// context is canceled
-				return
+			// wait until the slot is reached
+			for {
+				if ctx.Err() != nil {
+					// context is canceled
+					return
+				}
+
+				if wallet.CurrentSlot() == i {
+					break
+				}
+				time.Sleep(2 * time.Second)
 			}
 
-			d.IssueCandidacyPayloadFromAccount(wallet)
-			time.Sleep(time.Duration(slotDuration) * time.Second)
+			d.IssueCandidacyPayloadFromAccount(ctx, wallet)
 		}
 	}()
 }
@@ -318,7 +323,7 @@ func issueValidationBlockInBackground(ctx context.Context, wg *sync.WaitGroup, w
 					return
 				}
 
-				wallet.CreateAndSubmitValidationBlock(context.Background(), "", nil)
+				wallet.CreateAndSubmitValidationBlock(ctx, "", nil)
 				time.Sleep(1 * time.Second)
 			}
 		}