|
30 | 30 |
|
31 | 31 |
|
32 | 32 | _SYS_TESTS_DIR = os.path.realpath(os.path.dirname(__file__)) |
33 | | -LOGO_FILE = os.path.join(_SYS_TESTS_DIR, "data", "logo.png") |
34 | 33 | 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") |
39 | 36 | PROJECT_ID = os.environ.get("PROJECT_ID") |
40 | 37 |
|
41 | 38 |
|
@@ -180,6 +177,70 @@ def test_detect_logos_async(self): |
180 | 177 | assert logo_annotations[0]["description"] == "google" |
181 | 178 |
|
182 | 179 |
|
| 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 | + |
183 | 244 | @unittest.skipUnless(PROJECT_ID, "PROJECT_ID not set in environment.") |
184 | 245 | class TestVisionClientProductSearch(VisionSystemTestBase): |
185 | 246 | def setUp(self): |
@@ -541,7 +602,13 @@ def test_import_product_sets(self): |
541 | 602 | ) |
542 | 603 |
|
543 | 604 | # Verify the result. |
| 605 | + image_prefix = "import_sets_image_" |
544 | 606 | 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 | + ) |
546 | 613 | for status in response.result().statuses: |
547 | 614 | self.assertEqual(status.code, grpc.StatusCode.OK.value[0]) |
0 commit comments