fix: wait for load test to complete

workflows/testnet/testnet.go

@@ -246,64 +246,62 @@ func createWallets(ctx workflow.Context, req messages.TestnetWorkflowRequest, ch
 }
 
 func runLoadTest(ctx workflow.Context, req messages.TestnetWorkflowRequest, chainState, providerState []byte,
-	mnemonics []string, selector workflow.Selector) error {
+	mnemonics []string, selector workflow.Selector) (workflow.Future, error) {
 	if req.EthereumLoadTestSpec != nil {
-		workflow.Go(ctx, func(ctx workflow.Context) {
-			var loadTestResp messages.RunLoadTestResponse
-			f := workflow.ExecuteActivity(
-				workflow.WithStartToCloseTimeout(ctx, loadTestTimeout),
-				loadTestActivities.RunLoadTest,
-				messages.RunLoadTestRequest{
-					ChainState:      chainState,
-					ProviderState:   providerState,
-					LoadTestSpec:    *req.EthereumLoadTestSpec,
-					RunnerType:      req.RunnerType,
-					IsEvmChain:      req.IsEvmChain,
-					Mnemonics:       mnemonics,
-					CatalystVersion: req.CatalystVersion,
-				},
-			)
-
-			selector.AddFuture(f, func(f workflow.Future) {
-				activityErr := f.Get(ctx, &loadTestResp)
-				if activityErr != nil {
-					workflow.GetLogger(ctx).Error("ethereum load test failed", zap.Error(activityErr))
-				} else if loadTestResp.Result.Error != "" {
-					workflow.GetLogger(ctx).Error("ethereum load test reported an error", zap.String("error", loadTestResp.Result.Error))
-				}
-			})
-
+		var loadTestResp messages.RunLoadTestResponse
+		f := workflow.ExecuteActivity(
+			workflow.WithStartToCloseTimeout(ctx, loadTestTimeout),
+			loadTestActivities.RunLoadTest,
+			messages.RunLoadTestRequest{
+				ChainState:      chainState,
+				ProviderState:   providerState,
+				LoadTestSpec:    *req.EthereumLoadTestSpec,
+				RunnerType:      req.RunnerType,
+				IsEvmChain:      req.IsEvmChain,
+				Mnemonics:       mnemonics,
+				CatalystVersion: req.CatalystVersion,
+			},
+		)
+
+		selector.AddFuture(f, func(f workflow.Future) {
+			activityErr := f.Get(ctx, &loadTestResp)
+			if activityErr != nil {
+				workflow.GetLogger(ctx).Error("ethereum load test failed", zap.Error(activityErr))
+			} else if loadTestResp.Result.Error != "" {
+				workflow.GetLogger(ctx).Error("ethereum load test reported an error", zap.String("error", loadTestResp.Result.Error))
+			}
 		})
-	} else if req.CosmosLoadTestSpec != nil {
-		workflow.Go(ctx, func(ctx workflow.Context) {
-			var loadTestResp messages.RunLoadTestResponse
-			f := workflow.ExecuteActivity(
-				workflow.WithStartToCloseTimeout(ctx, loadTestTimeout),
-				loadTestActivities.RunLoadTest,
-				messages.RunLoadTestRequest{
-					ChainState:      chainState,
-					ProviderState:   providerState,
-					LoadTestSpec:    *req.CosmosLoadTestSpec,
-					RunnerType:      req.RunnerType,
-					IsEvmChain:      req.IsEvmChain,
-					Mnemonics:       mnemonics,
-					CatalystVersion: req.CatalystVersion,
-				},
-			)
-
-			selector.AddFuture(f, func(f workflow.Future) {
-				activityErr := f.Get(ctx, &loadTestResp)
-				if activityErr != nil {
-					workflow.GetLogger(ctx).Error("cosmos load test failed", zap.Error(activityErr))
-				} else if loadTestResp.Result.Error != "" {
-					workflow.GetLogger(ctx).Error("cosmos load test reported an error", zap.String("error", loadTestResp.Result.Error))
-				}
-			})
 
+		return f, nil
+	} else if req.CosmosLoadTestSpec != nil {
+		var loadTestResp messages.RunLoadTestResponse
+		f := workflow.ExecuteActivity(
+			workflow.WithStartToCloseTimeout(ctx, loadTestTimeout),
+			loadTestActivities.RunLoadTest,
+			messages.RunLoadTestRequest{
+				ChainState:      chainState,
+				ProviderState:   providerState,
+				LoadTestSpec:    *req.CosmosLoadTestSpec,
+				RunnerType:      req.RunnerType,
+				IsEvmChain:      req.IsEvmChain,
+				Mnemonics:       mnemonics,
+				CatalystVersion: req.CatalystVersion,
+			},
+		)
+
+		selector.AddFuture(f, func(f workflow.Future) {
+			activityErr := f.Get(ctx, &loadTestResp)
+			if activityErr != nil {
+				workflow.GetLogger(ctx).Error("cosmos load test failed", zap.Error(activityErr))
+			} else if loadTestResp.Result.Error != "" {
+				workflow.GetLogger(ctx).Error("cosmos load test reported an error", zap.String("error", loadTestResp.Result.Error))
+			}
 		})
+
+		return f, nil
 	}
 
-	return nil
+	return nil, nil
 }
 
 func startWorkflow(ctx workflow.Context, req messages.TestnetWorkflowRequest, runName string, buildResult messages.BuildDockerImageResponse, workflowID string) error {
@@ -334,7 +332,7 @@ func startWorkflow(ctx workflow.Context, req messages.TestnetWorkflowRequest, ru
 
 	shutdownSelector := workflow.NewSelector(ctx)
 	// 1. load test selector
-	err = runLoadTest(ctx, req, chainState, providerState, mnemonics, shutdownSelector)
+	loadTestFuture, err := runLoadTest(ctx, req, chainState, providerState, mnemonics, shutdownSelector)
 	if err != nil {
 		workflow.GetLogger(ctx).Error("load test initiation failed", zap.Error(err))
 	}
@@ -343,6 +341,18 @@ func startWorkflow(ctx workflow.Context, req messages.TestnetWorkflowRequest, ru
 
 	shutdownSelector.Select(ctx)
 
+	// If we have a loadtest running and the duration timer expired (not cancelled),
+	// wait for the loadtest to complete before allowing teardown
+	if loadTestFuture != nil && !temporal.IsCanceledError(ctx.Err()) {
+		workflow.GetLogger(ctx).Info("testnet duration expired but loadtest is still running, waiting for completion")
+		err = loadTestFuture.Get(ctx, nil) // Wait for loadtest to complete
+		if err != nil {
+			workflow.GetLogger(ctx).Error("failed to wait for load test future", zap.Error(err))
+			return err
+		}
+		workflow.GetLogger(ctx).Info("loadtest completed, proceeding with teardown")
+	}
+
 	if ctx.Err() != nil && temporal.IsCanceledError(ctx.Err()) {
 		workflow.GetLogger(ctx).Info("workflow was cancelled, completing gracefully")
 		return nil

workflows/testnet/testnet_test.go

@@ -553,6 +553,87 @@ func (s *TestnetWorkflowTestSuite) Test_TestnetWorkflowLongRunningCancelled() {
 	s.env.AssertActivityNumberOfCalls(s.T(), "TeardownProvider", 1)
 }
 
+func (s *TestnetWorkflowTestSuite) Test_TestnetWorkflowWaitsForLoadTestOnDurationExpiry() {
+	loadTestCompleted := false
+
+	cfg, err := types.ParseWorkerConfig("../../conf/worker.yaml")
+	if err != nil {
+		s.T().Fatal(err)
+	}
+	testnetActivity := &testnettypes.Activity{
+		Chains: cfg.Chains,
+	}
+	loadTestActivity := &loadtest.Activity{}
+	builderActivity := &builder.Activity{}
+	walletCreatorActivities := walletcreator.Activity{}
+
+	s.env.RegisterActivity(testnetActivity.CreateProvider)
+	s.env.RegisterActivity(testnetActivity.TeardownProvider)
+	s.env.RegisterActivity(testnetActivity.LaunchTestnet)
+	s.env.RegisterActivity(loadTestActivity.RunLoadTest)
+	s.env.RegisterActivity(builderActivity.BuildDockerImage)
+	s.env.RegisterActivity(walletCreatorActivities.CreateWallets)
+
+	testnetActivities = testnetActivity
+	loadTestActivities = loadTestActivity
+
+	s.env.OnActivity(loadTestActivity.RunLoadTest, mock.Anything, mock.Anything).Return(
+		func(ctx context.Context, req messages.RunLoadTestRequest) (messages.RunLoadTestResponse, error) {
+			time.Sleep(5 * time.Second)
+			loadTestCompleted = true
+			return messages.RunLoadTestResponse{
+				Result: catalysttypes.LoadTestResult{},
+			}, nil
+		})
+
+	s.env.OnActivity(testnetActivity.TeardownProvider, mock.Anything, mock.Anything).Return(
+		func(ctx context.Context, req messages.TeardownProviderRequest) (messages.TeardownProviderResponse, error) {
+			// fail if teardown is called before loadtest completes
+			if !loadTestCompleted {
+				s.T().Errorf("TeardownProvider called before LoadTest completed instead of waiting for load test to complete")
+			}
+			return messages.TeardownProviderResponse{}, nil
+		})
+
+	s.env.OnActivity(builderActivity.BuildDockerImage, mock.Anything, mock.Anything).Return(
+		func(ctx context.Context, req messages.BuildDockerImageRequest) (messages.BuildDockerImageResponse, error) {
+			originalTag := "ghcr.io/cosmos/simapp:v0.50"
+			newTag := "simapp-v53"
+
+			cmd := exec.Command("docker", "pull", originalTag)
+			output, err := cmd.CombinedOutput()
+			if err != nil {
+				return messages.BuildDockerImageResponse{}, err
+			}
+
+			tagCmd := exec.Command("docker", "tag", originalTag, newTag)
+			tagOutput, err := tagCmd.CombinedOutput()
+			if err != nil {
+				return messages.BuildDockerImageResponse{}, err
+			}
+
+			return messages.BuildDockerImageResponse{
+				FQDNTag: newTag,
+				Logs:    append(output, tagOutput...),
+			}, nil
+		})
+
+	dockerReq := simappReq
+	dockerReq.Repo = "cosmos-sdk"
+	dockerReq.SHA = "acb1d65cdc1e0fc36d93f3c5bb6aaf919a1321e2"
+	dockerReq.RunnerType = messages.Docker
+	dockerReq.ChainConfig.Name = "stake"
+	dockerReq.TestnetDuration = "3s"
+	dockerReq.LongRunningTestnet = false
+
+	s.env.ExecuteWorkflow(Workflow, dockerReq)
+
+	s.True(s.env.IsWorkflowCompleted())
+	s.NoError(s.env.GetWorkflowError())
+	s.env.AssertActivityNumberOfCalls(s.T(), "RunLoadTest", 1)
+	s.env.AssertActivityNumberOfCalls(s.T(), "TeardownProvider", 1)
+}
+
 func TestTestnetWorkflowTestSuite(t *testing.T) {
 	suite.Run(t, new(TestnetWorkflowTestSuite))
 }