Skip to content

Commit

Permalink
samples: added test that covers the wrong file type case (#69)
Browse files Browse the repository at this point in the history
* samples: added test that covers the wrong file type case
  • Loading branch information
munkhuushmgl authored and holtskinner committed Jan 3, 2023
1 parent 85ecf86 commit b1b0f92
Show file tree
Hide file tree
Showing 2 changed files with 47 additions and 1 deletion.
44 changes: 44 additions & 0 deletions batch_process_documents_sample_bad_input_v1beta3_test.py
Original file line number Diff line number Diff line change
@@ -0,0 +1,44 @@
# Copyright 2020 Google LLC
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
#

import os
from uuid import uuid4

from samples.snippets import batch_process_documents_sample_v1beta3

location = "us"
project_id = os.getenv("GOOGLE_CLOUD_PROJECT")
processor_id = "90484cfdedb024f6"
gcs_input_uri = "gs://cloud-samples-data/documentai/invoice.pdf"
# following bucket contains .csv file which will cause the sample to fail.
gcs_output_full_uri_with_wrong_type = "gs://documentai-beta-samples"
BUCKET_NAME = f"document-ai-python-{uuid4()}"


def test_batch_process_documents_with_bad_input(capsys):
try:
batch_process_documents_sample_v1beta3.batch_process_documents(
project_id=project_id,
location=location,
processor_id=processor_id,
gcs_input_uri=gcs_input_uri,
gcs_output_uri=gcs_output_full_uri_with_wrong_type,
gcs_output_uri_prefix="test",
timeout=450,
)
out, _ = capsys.readouterr()
assert "Failed to process" in out
except Exception as e:
assert "Failed to process" in e.message
4 changes: 3 additions & 1 deletion batch_process_documents_sample_v1beta3.py
Original file line number Diff line number Diff line change
Expand Up @@ -35,6 +35,7 @@ def batch_process_documents(
gcs_input_uri,
gcs_output_uri,
gcs_output_uri_prefix,
timeout: int = 300,
):

client = documentai.DocumentProcessorServiceClient()
Expand Down Expand Up @@ -63,7 +64,7 @@ def batch_process_documents(
operation = client.batch_process_documents(request)

# Wait for the operation to finish
operation.result()
operation.result(timeout=timeout)

# Results are written to GCS. Use a regex to find
# output files
Expand All @@ -79,6 +80,7 @@ def batch_process_documents(
for i, blob in enumerate(blob_list):
# Download the contents of this blob as a bytes object.
if ".json" not in blob.name:
print(f"skipping non-supported file type {blob.name}")
return
# Only parses JSON files
blob_as_bytes = blob.download_as_bytes()
Expand Down

0 comments on commit b1b0f92

Please sign in to comment.