try a retry mechanism

Author: yuandrew

temporal-spring-boot-autoconfigure/src/test/java/io/temporal/spring/boot/autoconfigure/WorkerVersioningTest.java

@@ -2,6 +2,8 @@
 
 import static org.junit.jupiter.api.Assertions.assertTrue;
 
+import io.grpc.Status;
+import io.grpc.StatusRuntimeException;
 import io.temporal.api.common.v1.WorkflowExecution;
 import io.temporal.api.enums.v1.EventType;
 import io.temporal.api.enums.v1.VersioningBehavior;
@@ -12,7 +14,13 @@
 import io.temporal.spring.boot.autoconfigure.workerversioning.TestWorkflow;
 import io.temporal.spring.boot.autoconfigure.workerversioning.TestWorkflow2;
 import io.temporal.worker.WorkerFactory;
-import org.junit.jupiter.api.*;
+import java.time.Duration;
+import org.junit.jupiter.api.Assumptions;
+import org.junit.jupiter.api.BeforeAll;
+import org.junit.jupiter.api.BeforeEach;
+import org.junit.jupiter.api.Test;
+import org.junit.jupiter.api.TestInstance;
+import org.junit.jupiter.api.Timeout;
 import org.springframework.beans.factory.annotation.Autowired;
 import org.springframework.boot.test.context.SpringBootTest;
 import org.springframework.context.ConfigurableApplicationContext;
@@ -50,15 +58,7 @@ public void testAutoDiscovery() {
     WorkerFactory workerFactory = applicationContext.getBean(WorkerFactory.class);
     workerFactory.start();
 
-    workflowClient
-        .getWorkflowServiceStubs()
-        .blockingStub()
-        .setWorkerDeploymentCurrentVersion(
-            SetWorkerDeploymentCurrentVersionRequest.newBuilder()
-                .setNamespace(workflowClient.getOptions().getNamespace())
-                .setDeploymentName("dname")
-                .setVersion("dname.bid")
-                .build());
+    setCurrentVersionWithRetry();
 
     TestWorkflow testWorkflow =
         workflowClient.newWorkflowStub(
@@ -91,6 +91,36 @@ public void testAutoDiscovery() {
                             == VersioningBehavior.VERSIONING_BEHAVIOR_AUTO_UPGRADE));
   }
 
+  @SuppressWarnings("deprecation")
+  private void setCurrentVersionWithRetry() {
+    long deadline = System.currentTimeMillis() + Duration.ofSeconds(10).toMillis();
+    while (true) {
+      try {
+        workflowClient
+            .getWorkflowServiceStubs()
+            .blockingStub()
+            .setWorkerDeploymentCurrentVersion(
+                SetWorkerDeploymentCurrentVersionRequest.newBuilder()
+                    .setNamespace(workflowClient.getOptions().getNamespace())
+                    .setDeploymentName("dname")
+                    .setVersion("dname.bid")
+                    .build());
+        return;
+      } catch (StatusRuntimeException e) {
+        if (e.getStatus().getCode() != Status.Code.NOT_FOUND
+            || System.currentTimeMillis() > deadline) {
+          throw e;
+        }
+        try {
+          Thread.sleep(100);
+        } catch (InterruptedException ie) {
+          Thread.currentThread().interrupt();
+          throw new RuntimeException(ie);
+        }
+      }
+    }
+  }
+
   @ComponentScan(
       excludeFilters =
           @ComponentScan.Filter(