chore: added increased timeout on flaky batch request (#84)

GoogleCloudPlatform · Jan 3, 2023 · 87b9220 · 87b9220
1 parent 2774619
commit 87b9220
Show file tree

Hide file tree

Showing 4 changed files with 6 additions and 4 deletions.
diff --git a/batch_parse_form_v1beta2.py b/batch_parse_form_v1beta2.py
@@ -24,6 +24,7 @@ def batch_parse_form(
     project_id="YOUR_PROJECT_ID",
     input_uri="gs://cloud-samples-data/documentai/form.pdf",
     destination_uri="gs://your-bucket-id/path/to/save/results/",
+    timeout=90
 ):
     """Parse a form"""
 
@@ -80,7 +81,7 @@ def batch_parse_form(
     operation = client.batch_process_documents(batch_request)
 
     # Wait for the operation to finish
-    operation.result()
+    operation.result(timeout)
 
     # Results are written to GCS. Use a regex to find
     # output files

diff --git a/batch_parse_form_v1beta2_test.py b/batch_parse_form_v1beta2_test.py
@@ -41,6 +41,6 @@ def setup_teardown():
 
 
 def test_batch_parse_form(capsys):
-    batch_parse_form_v1beta2.batch_parse_form(PROJECT_ID, INPUT_URI, BATCH_OUTPUT_URI)
+    batch_parse_form_v1beta2.batch_parse_form(PROJECT_ID, INPUT_URI, BATCH_OUTPUT_URI, 120)
     out, _ = capsys.readouterr()
     assert "Output files" in out
diff --git a/batch_parse_table_v1beta2.py b/batch_parse_table_v1beta2.py
@@ -24,6 +24,7 @@ def batch_parse_table(
     project_id="YOUR_PROJECT_ID",
     input_uri="gs://cloud-samples-data/documentai/form.pdf",
     destination_uri="gs://your-bucket-id/path/to/save/results/",
+    timeout=90
 ):
     """Parse a form"""
 
@@ -88,7 +89,7 @@ def batch_parse_table(
     operation = client.batch_process_documents(batch_request)
 
     # Wait for the operation to finish
-    operation.result()
+    operation.result(timeout)
 
     # Results are written to GCS. Use a regex to find
     # output files

diff --git a/batch_parse_table_v1beta2_test.py b/batch_parse_table_v1beta2_test.py
@@ -41,6 +41,6 @@ def setup_teardown():
 
 
 def test_batch_parse_table(capsys):
-    batch_parse_table_v1beta2.batch_parse_table(PROJECT_ID, INPUT_URI, BATCH_OUTPUT_URI)
+    batch_parse_table_v1beta2.batch_parse_table(PROJECT_ID, INPUT_URI, BATCH_OUTPUT_URI, 120)
     out, _ = capsys.readouterr()
     assert "Output files:" in out