From eb5a45475c3c9b3dea9aa8671f9f5377b645cb76 Mon Sep 17 00:00:00 2001
From: Tanmay <tanmay@zetachain.com>
Date: Mon, 30 Sep 2024 15:55:01 +0530
Subject: [PATCH] test: add wait for block to tss migration test (#2931)

* add wait for block to tss migration test

* add comments

* refactor identifiers

* rename checkNumberOfTssGenerated to checkNumberOfTSSGenerated
---
 e2e/e2etests/test_migrate_tss.go |  2 ++
 e2e/runner/zeta.go               | 29 +++++++++++++++++++++++++++++
 2 files changed, 31 insertions(+)

diff --git a/e2e/e2etests/test_migrate_tss.go b/e2e/e2etests/test_migrate_tss.go
index cb5e3c35c4..3825228c14 100644
--- a/e2e/e2etests/test_migrate_tss.go
+++ b/e2e/e2etests/test_migrate_tss.go
@@ -48,6 +48,8 @@ func TestMigrateTSS(r *runner.E2ERunner, _ []string) {
 	btcBalance = btcBalance - 0.01
 	btcChain := chains.BitcoinRegtest.ChainId
 
+	r.WaitForTSSGeneration(2)
+
 	//migrate btc funds
 	// #nosec G701 e2eTest - always in range
 	migrationAmountBTC := sdkmath.NewUint(uint64(btcBalance * 1e8))
diff --git a/e2e/runner/zeta.go b/e2e/runner/zeta.go
index cec59219cd..bdada4763c 100644
--- a/e2e/runner/zeta.go
+++ b/e2e/runner/zeta.go
@@ -15,8 +15,11 @@ import (
 	"github.com/zeta-chain/node/e2e/utils"
 	"github.com/zeta-chain/node/pkg/retry"
 	"github.com/zeta-chain/node/x/crosschain/types"
+	observertypes "github.com/zeta-chain/node/x/observer/types"
 )
 
+// WaitForBlocks waits for a specific number of blocks to be generated
+// The parameter n is the number of blocks to wait for
 func (r *E2ERunner) WaitForBlocks(n int64) {
 	height, err := r.CctxClient.LastZetaHeight(r.Ctx, &types.QueryLastZetaHeightRequest{})
 	if err != nil {
@@ -31,6 +34,32 @@ func (r *E2ERunner) WaitForBlocks(n int64) {
 	err = retry.DoWithBackoff(call, boWithMaxRetries)
 	require.NoError(r, err, "failed to wait for %d blocks", n)
 }
+
+// WaitForTSSGeneration waits for a specific number of TSS to be generated
+// The parameter n is the number of TSS to wait for
+func (r *E2ERunner) WaitForTSSGeneration(tssNumber int64) {
+	call := func() error {
+		return retry.Retry(r.checkNumberOfTSSGenerated(tssNumber))
+	}
+	bo := backoff.NewConstantBackOff(time.Second * 5)
+	boWithMaxRetries := backoff.WithMaxRetries(bo, 10)
+	err := retry.DoWithBackoff(call, boWithMaxRetries)
+	require.NoError(r, err, "failed to wait for %d tss generation", tssNumber)
+}
+
+// checkNumberOfTSSGenerated checks the number of TSS generated
+// if the number of tss is less that the `tssNumber` provided we return an error
+func (r *E2ERunner) checkNumberOfTSSGenerated(tssNumber int64) error {
+	tssList, err := r.ObserverClient.TssHistory(r.Ctx, &observertypes.QueryTssHistoryRequest{})
+	if err != nil {
+		return err
+	}
+	if int64(len(tssList.TssList)) < tssNumber {
+		return fmt.Errorf("waiting for %d tss generation, number of TSS :%d", tssNumber, len(tssList.TssList))
+	}
+	return nil
+}
+
 func (r *E2ERunner) waitForBlock(n int64) error {
 	height, err := r.CctxClient.LastZetaHeight(r.Ctx, &types.QueryLastZetaHeightRequest{})
 	if err != nil {