Skip to content

Commit 879ef99

Browse files
brendanlundybusunkim96
authored andcommitted
Add AsyncBatchAnnotateFiles system test (googleapis#8260)
1 parent 42ffb49 commit 879ef99

6 files changed

Lines changed: 73 additions & 6 deletions

File tree

vision/tests/data/car.jpg

-65.5 KB
Binary file not shown.

vision/tests/data/full-text.jpg

-1.11 MB
Binary file not shown.

vision/tests/data/landmark.jpg

-92.9 KB
Binary file not shown.

vision/tests/data/pdf_test.pdf

15.2 KB
Binary file not shown.

vision/tests/data/text.jpg

-23.3 KB
Binary file not shown.

vision/tests/system.py

Lines changed: 73 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -30,12 +30,9 @@
3030

3131

3232
_SYS_TESTS_DIR = os.path.realpath(os.path.dirname(__file__))
33-
LOGO_FILE = os.path.join(_SYS_TESTS_DIR, "data", "logo.png")
3433
FACE_FILE = os.path.join(_SYS_TESTS_DIR, "data", "faces.jpg")
35-
LABEL_FILE = os.path.join(_SYS_TESTS_DIR, "data", "car.jpg")
36-
LANDMARK_FILE = os.path.join(_SYS_TESTS_DIR, "data", "landmark.jpg")
37-
TEXT_FILE = os.path.join(_SYS_TESTS_DIR, "data", "text.jpg")
38-
FULL_TEXT_FILE = os.path.join(_SYS_TESTS_DIR, "data", "full-text.jpg")
34+
LOGO_FILE = os.path.join(_SYS_TESTS_DIR, "data", "logo.png")
35+
PDF_FILE = os.path.join(_SYS_TESTS_DIR, "data", "pdf_test.pdf")
3936
PROJECT_ID = os.environ.get("PROJECT_ID")
4037

4138

@@ -180,6 +177,70 @@ def test_detect_logos_async(self):
180177
assert logo_annotations[0]["description"] == "google"
181178

182179

180+
class TestVisionClientFiles(VisionSystemTestBase):
181+
def test_async_batch_annotate_files(self):
182+
# Upload the image to Google Cloud Storage.
183+
blob_name = "async_batch_annotate_files.pdf"
184+
blob = self.test_bucket.blob(blob_name)
185+
self.to_delete_by_case.append(blob)
186+
with io.open(PDF_FILE, "rb") as image_file:
187+
blob.upload_from_file(image_file)
188+
189+
# Make the request.
190+
method_name = "test_async_batch_annotate_files"
191+
output_gcs_uri_prefix = "gs://{bucket}/{method_name}".format(
192+
bucket=self.test_bucket.name, method_name=method_name
193+
)
194+
request = {
195+
"input_config": {
196+
"gcs_source": {
197+
"uri": "gs://{bucket}/{blob}".format(
198+
bucket=self.test_bucket.name, blob=blob_name
199+
)
200+
},
201+
"mime_type": "application/pdf",
202+
},
203+
"features": [{"type": vision.enums.Feature.Type.DOCUMENT_TEXT_DETECTION}],
204+
"output_config": {"gcs_destination": {"uri": output_gcs_uri_prefix}},
205+
}
206+
response = self.client.async_batch_annotate_files([request])
207+
208+
# Wait for the operation to complete.
209+
lro_waiting_seconds = 60
210+
start_time = time.time()
211+
while not response.done() and (time.time() - start_time) < lro_waiting_seconds:
212+
time.sleep(1)
213+
214+
if not response.done():
215+
self.fail(
216+
"{method_name} timed out after {lro_waiting_seconds} seconds".format(
217+
method_name=method_name, lro_waiting_seconds=lro_waiting_seconds
218+
)
219+
)
220+
221+
# Make sure getting the result is not an error.
222+
response.result()
223+
224+
# There should be exactly 1 output file in gcs at the prefix output_gcs_uri_prefix.
225+
blobs = list(self.test_bucket.list_blobs(prefix=method_name))
226+
assert len(blobs) == 1
227+
blob = blobs[0]
228+
229+
# Download the output file and verify the result
230+
result_str = blob.download_as_string().decode("utf8")
231+
result = json.loads(result_str)
232+
responses = result["responses"]
233+
assert len(responses) == 1
234+
text = responses[0]["fullTextAnnotation"]["text"]
235+
expected_text = "test text"
236+
self.assertTrue(
237+
expected_text in text,
238+
"'{expected_text}' not in '{text}'".format(
239+
expected_text=expected_text, text=text
240+
),
241+
)
242+
243+
183244
@unittest.skipUnless(PROJECT_ID, "PROJECT_ID not set in environment.")
184245
class TestVisionClientProductSearch(VisionSystemTestBase):
185246
def setUp(self):
@@ -541,7 +602,13 @@ def test_import_product_sets(self):
541602
)
542603

543604
# Verify the result.
605+
image_prefix = "import_sets_image_"
544606
for ref_image in response.result().reference_images:
545-
self.assertTrue("import_sets_image_" in ref_image.uri)
607+
self.assertTrue(
608+
image_prefix in ref_image.uri,
609+
"'{image_prefix}' not in '{uri}'".format(
610+
image_prefix=image_prefix, uri=ref_image.uri
611+
),
612+
)
546613
for status in response.result().statuses:
547614
self.assertEqual(status.code, grpc.StatusCode.OK.value[0])

0 commit comments

Comments
 (0)