fix: simplify duplicate retry operations

Signed-off-by: mao3267 <chenvincent610@gmail.com>
flyteorg · Jan 11, 2025 · f167305 · f167305
1 parent 58c2e91
commit f167305
Showing 1 changed file with 12 additions and 29 deletions.
diff --git a/tests/flytekit/integration/remote/test_remote.py b/tests/flytekit/integration/remote/test_remote.py
@@ -876,43 +876,26 @@ def test_signal_approve_reject(register):
 
     execution = remote.execute(conditional_wf, inputs={"data": [1.0, 2.0, 3.0, 4.0, 5.0]})
 
-    max_retries = 10
+    def retry_operation(operation):
+        max_retries = 10
+        for _ in range(max_retries):
+            try:
+                operation()
+                break
+            except Exception:
+                sleep(1)
 
-    for _ in range(max_retries):
-        try:
-            remote.set_input("title-input", execution.id.name, value="my report", project=PROJECT, domain=DOMAIN, python_type=str, literal_type=LiteralType(simple=SimpleType.STRING))
-            break
-        except Exception as e:
-            sleep(1)
-            continue
-    for _ in range(max_retries):
-        try:
-            remote.approve("review-passes", execution.id.name, project=PROJECT, domain=DOMAIN)
-            break
-        except Exception as e:
-            sleep(1)
-            continue
+    retry_operation(lambda: remote.set_input("title-input", execution.id.name, value="my report", project=PROJECT, domain=DOMAIN, python_type=str, literal_type=LiteralType(simple=SimpleType.STRING)))
+    retry_operation(lambda: remote.approve("review-passes", execution.id.name, project=PROJECT, domain=DOMAIN))
 
     remote.wait(execution=execution, timeout=datetime.timedelta(minutes=5))
     assert execution.outputs["o0"] == {"title": "my report", "data": [1.0, 2.0, 3.0, 4.0, 5.0]}
 
     with pytest.raises(FlyteAssertion, match="Outputs could not be found because the execution ended in failure"):
         execution = remote.execute(conditional_wf, inputs={"data": [1.0, 2.0, 3.0, 4.0, 5.0]})
 
-        for _ in range(max_retries):
-            try:
-                remote.set_input("title-input", execution.id.name, value="my report", project=PROJECT, domain=DOMAIN, python_type=str, literal_type=LiteralType(simple=SimpleType.STRING))
-                break
-            except Exception as e:
-                sleep(1)
-                continue
-        for _ in range(max_retries):
-            try:
-                remote.reject("review-passes", execution.id.name, project=PROJECT, domain=DOMAIN)
-                break
-            except Exception as e:
-                sleep(1)
-                continue
+        retry_operation(lambda: remote.set_input("title-input", execution.id.name, value="my report", project=PROJECT, domain=DOMAIN, python_type=str, literal_type=LiteralType(simple=SimpleType.STRING)))
+        retry_operation(lambda: remote.reject("review-passes", execution.id.name, project=PROJECT, domain=DOMAIN))
 
         remote.wait(execution=execution, timeout=datetime.timedelta(minutes=5))
         assert execution.outputs["o0"] == {"title": "my report", "data": [1.0, 2.0, 3.0, 4.0, 5.0]}