From 87b922088802d8756ba39be7778047d00cdac9da Mon Sep 17 00:00:00 2001
From: Mike <45373284+munkhuushmgl@users.noreply.github.com>
Date: Fri, 22 Jan 2021 10:45:09 -0800
Subject: [PATCH] chore: added increased timeout on flaky batch request (#84)

---
 batch_parse_form_v1beta2.py       | 3 ++-
 batch_parse_form_v1beta2_test.py  | 2 +-
 batch_parse_table_v1beta2.py      | 3 ++-
 batch_parse_table_v1beta2_test.py | 2 +-
 4 files changed, 6 insertions(+), 4 deletions(-)

diff --git a/batch_parse_form_v1beta2.py b/batch_parse_form_v1beta2.py
index 01c19e1e716c..ae60fd6334ba 100644
--- 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
index 50dc845d4ea6..6abd19a20559 100644
--- 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
index 08942080684a..f62495b4f340 100644
--- 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
index ed1be2ee070f..aa8905208500 100644
--- 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