Add ImportProductSets system test (#8259)

packages/google-cloud-vision/tests/system.py

@@ -14,6 +14,7 @@
 
 """System tests for Vision API."""
 
+import grpc
 import io
 import json
 import os
@@ -200,6 +201,15 @@ def tearDown(self):
         for product_set in self.product_sets_to_delete:
             self.ps_client.delete_product_set(name=product_set)
 
+    def _upload_image(self, image_name):
+        blob = self.test_bucket.blob(image_name)
+        self.to_delete_by_case.append(blob)
+        with io.open(FACE_FILE, "rb") as image_file:
+            blob.upload_from_file(image_file)
+        return "gs://{bucket}/{blob}".format(
+            bucket=self.test_bucket.name, blob=image_name
+        )
+
     def test_create_product_set(self):
         # Create a ProductSet.
         product_set = vision.types.ProductSet(display_name="display name")
@@ -420,15 +430,10 @@ def test_reference_image(self):
         )
         self.products_to_delete.append(response.name)
         self.assertEqual(response.name, product_path)
+
         # Upload image to gcs.
-        blob_name = "faces.jpg"
-        blob = self.test_bucket.blob(blob_name)
-        self.to_delete_by_case.append(blob)
-        with io.open(FACE_FILE, "rb") as image_file:
-            blob.upload_from_file(image_file)
-        gcs_uri = "gs://{bucket}/{blob}".format(
-            bucket=self.test_bucket.name, blob=blob_name
-        )
+        gcs_uri = self._upload_image("reference_image_test.jpg")
+
         # Create a ReferenceImage.
         reference_image_id = "reference_image" + unique_resource_id()
         reference_image_path = self.ps_client.reference_image_path(
@@ -445,12 +450,98 @@ def test_reference_image(self):
         )
         self.reference_images_to_delete.append(response.name)
         self.assertEqual(response.name, reference_image_path)
+
         # Get the ReferenceImage.
         get_response = self.ps_client.get_reference_image(name=reference_image_path)
         self.assertEqual(get_response.name, reference_image_path)
+
         # List the ReferenceImages in the Product.
         listed_reference_images = list(
             self.ps_client.list_reference_images(parent=product_path)
         )
         self.assertEqual(len(listed_reference_images), 1)
         self.assertEqual(listed_reference_images[0].name, reference_image_path)
+
+    def _build_csv_line(
+        self, gcs_uri_image, reference_image_id, product_set_id, product_id
+    ):
+        return ",".join(
+            [
+                gcs_uri_image,
+                reference_image_id,
+                product_set_id,
+                product_id,
+                "apparel",
+                "display name",
+                '"color=black,style=formal"',
+                "",
+            ]
+        )
+
+    def test_import_product_sets(self):
+        # Generate the ids that will be used in the import.
+        product_set_id = "set" + unique_resource_id()
+        product_set_path = self.ps_client.product_set_path(
+            project=PROJECT_ID, location=self.location, product_set=product_set_id
+        )
+        self.product_sets_to_delete.append(product_set_path)
+        product_id = "product" + unique_resource_id()
+        product_path = self.ps_client.product_path(
+            project=PROJECT_ID, location=self.location, product=product_id
+        )
+        self.products_to_delete.append(product_path)
+        reference_image_id_1 = "reference_image_1" + unique_resource_id()
+        reference_image_path = self.ps_client.reference_image_path(
+            project=PROJECT_ID,
+            location=self.location,
+            product=product_id,
+            reference_image=reference_image_id_1,
+        )
+        self.reference_images_to_delete.append(reference_image_path)
+        reference_image_id_2 = "reference_image_2" + unique_resource_id()
+        reference_image_path = self.ps_client.reference_image_path(
+            project=PROJECT_ID,
+            location=self.location,
+            product=product_id,
+            reference_image=reference_image_id_2,
+        )
+        self.reference_images_to_delete.append(reference_image_path)
+
+        # Upload images to gcs.
+        gcs_uri_image_1 = self._upload_image("import_sets_image_1.jpg")
+        gcs_uri_image_2 = self._upload_image("import_sets_image_2.jpg")
+
+        # Build the string that will be uploaded to gcs as a csv file.
+        csv_data = "\n".join(
+            [
+                self._build_csv_line(
+                    gcs_uri_image_1, reference_image_id_1, product_set_id, product_id
+                ),
+                self._build_csv_line(
+                    gcs_uri_image_2, reference_image_id_2, product_set_id, product_id
+                ),
+            ]
+        )
+
+        # Upload a csv file to gcs.
+        csv_filename = "import_sets.csv"
+        blob = self.test_bucket.blob(csv_filename)
+        self.to_delete_by_case.append(blob)
+        blob.upload_from_string(csv_data)
+
+        # Make the import_product_sets request.
+        gcs_source = vision.types.ImportProductSetsGcsSource(
+            csv_file_uri="gs://{bucket}/{blob}".format(
+                bucket=self.test_bucket.name, blob=csv_filename
+            )
+        )
+        input_config = vision.types.ImportProductSetsInputConfig(gcs_source=gcs_source)
+        response = self.ps_client.import_product_sets(
+            parent=self.location_path, input_config=input_config
+        )
+
+        # Verify the result.
+        for ref_image in response.result().reference_images:
+            self.assertTrue("import_sets_image_" in ref_image.uri)
+        for status in response.result().statuses:
+            self.assertEqual(status.code, grpc.StatusCode.OK.value[0])