From 4dca5c05ee1e37b8da8457f57cad104b7b12ceca Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Philip=20Dub=C3=A9?= <philip@peerdb.io>
Date: Thu, 29 Feb 2024 02:35:24 +0000
Subject: [PATCH] WithTimeout to cap time spent asking to describe

---
 flow/e2e/test_utils.go | 6 ++++--
 1 file changed, 4 insertions(+), 2 deletions(-)

diff --git a/flow/e2e/test_utils.go b/flow/e2e/test_utils.go
index 57a624c1f0..251177bcb2 100644
--- a/flow/e2e/test_utils.go
+++ b/flow/e2e/test_utils.go
@@ -143,7 +143,7 @@ func EnvWaitForEqualTablesWithNames(
 
 func RequireEnvCanceled(t *testing.T, env WorkflowRun) {
 	t.Helper()
-	EnvWaitFor(t, env, time.Minute, "canceling", env.Finished)
+	EnvWaitForFinished(t, env, time.Minute)
 	err := env.Error()
 	var panicErr *temporal.PanicError
 	var canceledErr *temporal.CanceledError
@@ -637,7 +637,9 @@ func EnvWaitForFinished(t *testing.T, env WorkflowRun, timeout time.Duration) {
 	t.Helper()
 
 	EnvWaitFor(t, env, timeout, "finish", func() bool {
-		desc, err := env.c.DescribeWorkflowExecution(context.Background(), env.GetID(), "")
+		ctx, cancel := context.WithTimeout(context.Background(), 10*time.Second)
+		defer cancel()
+		desc, err := env.c.DescribeWorkflowExecution(ctx, env.GetID(), "")
 		if err != nil {
 			t.Log("Not finished", err)
 			return false