Update GriptapeCloudStructureRunDriver to look for completed statuses

griptape-ai · Sep 17, 2024 · ea9039b · ea9039b
1 parent c3f5d49
commit ea9039b
Show file tree

Hide file tree

Showing 3 changed files with 14 additions and 7 deletions.
diff --git a/griptape/drivers/structure_run/griptape_cloud_structure_run_driver.py b/griptape/drivers/structure_run/griptape_cloud_structure_run_driver.py
@@ -50,7 +50,10 @@ def _get_structure_run_result(self, structure_run_id: str) -> BaseArtifact | Inf
         status = result["status"]
 
         wait_attempts = 0
-        while status in ("QUEUED", "RUNNING") and wait_attempts < self.structure_run_max_wait_time_attempts:
+        while (
+            status not in ("SUCCEEDED", "FAILED", "ERROR", "CANCELLED")
+            and wait_attempts < self.structure_run_max_wait_time_attempts
+        ):
             # wait
             time.sleep(self.structure_run_wait_time_interval)
             wait_attempts += 1

diff --git a/tests/integration/test_code_blocks.py b/tests/integration/test_code_blocks.py
@@ -19,6 +19,7 @@
     "docs/griptape-framework/drivers/src/observability_drivers_2.py",
     "docs/griptape-framework/structures/src/observability_1.py",
     "docs/griptape-framework/structures/src/observability_2.py",
+    "docs/griptape-framework/data/src/loaders_9.py",
 ]
 
 

diff --git a/tests/unit/drivers/structure_run/test_griptape_cloud_structure_run_driver.py b/tests/unit/drivers/structure_run/test_griptape_cloud_structure_run_driver.py
@@ -15,15 +15,18 @@ def driver(self, mocker, mock_requests_post):
         from griptape.drivers import GriptapeCloudStructureRunDriver
 
         mock_response = mocker.Mock()
-        mock_response.json.return_value = {
-            "description": "fizz buzz",
-            "output": TextArtifact("foo bar").to_dict(),
-            "status": "SUCCEEDED",
-        }
+        mock_response.json.side_effect = [
+            {"description": "fizz buzz", "status": "RUNNING"},
+            {"description": "fizz buzz", "output": TextArtifact("foo bar").to_dict(), "status": "SUCCEEDED"},
+        ]
         mocker.patch("requests.get", return_value=mock_response)
 
         return GriptapeCloudStructureRunDriver(
-            base_url="https://cloud-foo.griptape.ai", api_key="foo bar", structure_id="1", env={"key": "value"}
+            base_url="https://cloud-foo.griptape.ai",
+            api_key="foo bar",
+            structure_id="1",
+            env={"key": "value"},
+            structure_run_wait_time_interval=0,
         )
 
     def test_run(self, driver, mock_requests_post):