From c11716660bfee70506a7aaad7a01e9e71b748417 Mon Sep 17 00:00:00 2001
From: Matt Vallillo <matt@griptape.ai>
Date: Tue, 17 Sep 2024 14:04:34 -0400
Subject: [PATCH] Update `GriptapeCloudStructureRunDriver` to look for
 completed statuses (#1180)

---
 .../griptape_cloud_structure_run_driver.py        |  5 ++++-
 tests/integration/test_code_blocks.py             |  1 +
 .../test_griptape_cloud_structure_run_driver.py   | 15 +++++++++------
 3 files changed, 14 insertions(+), 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
index 46b54c528..cc83797d7 100644
--- 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
index c9c666b0e..409bdd599 100644
--- 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
index ccc8ac303..bbd6bac2b 100644
--- 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):