def test_update_product_no_explicit_name(self, get_conn):
# Given
product = Product()
update_product_method = get_conn.return_value.update_product
update_product_method.return_value = product
product_name = ProductSearchClient.product_path(
PROJECT_ID_TEST, LOC_ID_TEST, PRODUCT_ID_TEST)
# When
result = self.hook.update_product(
location=LOC_ID_TEST,
product_id=PRODUCT_ID_TEST,
product=product,
update_mask=None,
project_id=PROJECT_ID_TEST,
retry=None,
timeout=None,
metadata=None,
)
# Then
self.assertEqual(result, MessageToDict(product))
update_product_method.assert_called_once_with(
product=Product(name=product_name),
metadata=None,
retry=None,
timeout=None,
update_mask=None)
def test_update_product_explicit_name_missing_params_for_constructed_name(
self, location, product_id, get_conn
):
# Given
explicit_p_name = ProductSearchClient.product_path(
PROJECT_ID_TEST_2, LOC_ID_TEST_2, PRODUCT_ID_TEST_2
)
product = Product(name=explicit_p_name)
update_product_method = get_conn.return_value.update_product
update_product_method.return_value = product
# When
result = self.hook.update_product(
location=location,
product_id=product_id,
product=product,
update_mask=None,
project_id=PROJECT_ID_TEST,
retry=None,
timeout=None,
metadata=None,
)
# Then
self.assertEqual(result, MessageToDict(product))
update_product_method.assert_called_once_with(
product=Product(name=explicit_p_name), metadata=None, retry=None, timeout=None, update_mask=None
)
def test_create_product_autogenerated_id_wrong_name_in_response(
self, get_conn):
# Given
wrong_name = 'wrong_name_not_a_correct_path'
response_product = Product(name=wrong_name)
create_product_method = get_conn.return_value.create_product
create_product_method.return_value = response_product
parent = ProductSearchClient.location_path(PROJECT_ID_TEST,
LOC_ID_TEST)
product = Product()
# When
with self.assertRaises(AirflowException) as cm:
self.hook.create_product(location=LOC_ID_TEST,
product_id=None,
product=product,
project_id=PROJECT_ID_TEST)
# Then
# API response was wrong (wrong name format) and thus ProductSet ID extraction should fail.
err = cm.exception
self.assertIn('Unable to get id from name', str(err))
create_product_method.assert_called_once_with(parent=parent,
product=product,
product_id=None,
retry=None,
timeout=None,
metadata=None)
def test_create_product_autogenerated_id(self, get_conn):
# Given
autogenerated_id = 'autogen-p-id'
response_product = Product(name=ProductSearchClient.product_path(
PROJECT_ID_TEST, LOC_ID_TEST, autogenerated_id))
create_product_method = get_conn.return_value.create_product
create_product_method.return_value = response_product
parent = ProductSearchClient.location_path(PROJECT_ID_TEST,
LOC_ID_TEST)
product = Product()
# When
result = self.hook.create_product(location=LOC_ID_TEST,
product_id=None,
product=product,
project_id=PROJECT_ID_TEST)
# Then
# Product ID was not provided in the method call above. Should be extracted from the API response
# and returned.
self.assertEqual(result, autogenerated_id)
create_product_method.assert_called_once_with(parent=parent,
product=product,
product_id=None,
retry=None,
timeout=None,
metadata=None)
def test_update_product_explicit_name_different_from_constructed(self, get_conn):
# Given
update_product_method = get_conn.return_value.update_product
update_product_method.return_value = None
explicit_p_name = ProductSearchClient.product_path(
PROJECT_ID_TEST_2, LOC_ID_TEST_2, PRODUCT_ID_TEST_2
)
product = Product(name=explicit_p_name)
template_p_name = ProductSearchClient.product_path(PROJECT_ID_TEST, LOC_ID_TEST, PRODUCT_ID_TEST)
# When
# Location and product_id are passed in addition to a Product with an explicit name,
# but both names differ (constructed != explicit).
# Should throw AirflowException in this case.
with self.assertRaises(AirflowException) as cm:
self.hook.update_product(
location=LOC_ID_TEST,
product_id=PRODUCT_ID_TEST,
product=product,
update_mask=None,
project_id=PROJECT_ID_TEST,
retry=None,
timeout=None,
metadata=None,
)
err = cm.exception
self.assertTrue(err)
self.assertIn(
"The Product name provided in the object ({}) is different than the name created from the input "
"parameters ({}). Please either: 1) Remove the Product name, 2) Remove the location and product_"
"id parameters, 3) Unify the Product name and input parameters.".format(
explicit_p_name, template_p_name
),
str(err),
)
update_product_method.assert_not_called()
def test_update_product_no_explicit_name_and_missing_params_for_constructed_name(
self, location, product_id, get_conn
):
# Given
update_product_method = get_conn.return_value.update_product
update_product_method.return_value = None
product = Product()
# When
with self.assertRaises(AirflowException) as cm:
self.hook.update_product(
location=location,
product_id=product_id,
product=product,
update_mask=None,
project_id=PROJECT_ID_TEST,
retry=None,
timeout=None,
metadata=None,
)
err = cm.exception
self.assertTrue(err)
self.assertIn(
"Unable to determine the Product name. Please either set the name directly in the "
"Product object or provide the `location` and `product_id` parameters.",
str(err),
)
update_product_method.assert_not_called()
def test_update_product_explicit_name_different_from_constructed(
self, get_conn):
# Given
update_product_method = get_conn.return_value.update_product
update_product_method.return_value = None
explicit_p_name = ProductSearchClient.product_path(
PROJECT_ID_TEST_2, LOC_ID_TEST_2, PRODUCT_ID_TEST_2)
product = Product(name=explicit_p_name)
template_p_name = ProductSearchClient.product_path(
PROJECT_ID_TEST, LOC_ID_TEST, PRODUCT_ID_TEST)
# When
# Location and product_id are passed in addition to a Product with an explicit name,
# but both names differ (constructed != explicit).
# Should throw AirflowException in this case.
with self.assertRaises(AirflowException) as cm:
self.hook.update_product(
location=LOC_ID_TEST,
product_id=PRODUCT_ID_TEST,
product=product,
update_mask=None,
project_id=PROJECT_ID_TEST,
retry=None,
timeout=None,
metadata=None,
)
err = cm.exception
self.assertTrue(err)
self.assertIn(
ERR_DIFF_NAMES.format(explicit_name=explicit_p_name,
constructed_name=template_p_name,
label="Product",
id_label="product_id"),
str(err),
)
update_product_method.assert_not_called()
def test_update_product_no_explicit_name_and_missing_params_for_constructed_name(
self, location, product_id, get_conn):
# Given
update_product_method = get_conn.return_value.update_product
update_product_method.return_value = None
product = Product()
# When
with self.assertRaises(AirflowException) as cm:
self.hook.update_product(
location=location,
product_id=product_id,
product=product,
update_mask=None,
project_id=PROJECT_ID_TEST,
retry=None,
timeout=None,
metadata=None,
)
err = cm.exception
self.assertTrue(err)
self.assertIn(
ERR_UNABLE_TO_CREATE.format(label='Product',
id_label='product_id'),
str(err),
)
update_product_method.assert_not_called()
def test_create_product_autogenerated_id_wrong_api_response(
self, get_conn):
# Given
response_product = None
create_product_method = get_conn.return_value.create_product
create_product_method.return_value = response_product
parent = ProductSearchClient.location_path(PROJECT_ID_TEST,
LOC_ID_TEST)
hook = self.vision_hook_default_project_id
product = Product()
# When
with self.assertRaises(AirflowException) as cm:
hook.create_product(location=LOC_ID_TEST,
product_id=None,
product=product,
project_id=PROJECT_ID_TEST)
# Then
# API response was wrong (None) and thus ProductSet ID extraction should fail.
err = cm.exception
self.assertIn('Unable to get name from response...', str(err))
create_product_method.assert_called_once_with(parent=parent,
product=product,
product_id=None,
retry=None,
timeout=None,
metadata=None)
def test_create_product_explicit_id(self, get_conn):
# Given
create_product_method = get_conn.return_value.create_product
create_product_method.return_value = None
parent = ProductSearchClient.location_path(PROJECT_ID_TEST, LOC_ID_TEST)
product = Product()
# When
result = self.hook.create_product(
location=LOC_ID_TEST, product_id=PRODUCT_ID_TEST, product=product, project_id=PROJECT_ID_TEST
)
# Then
# Product ID was provided explicitly in the method call above, should be returned from the method
self.assertEqual(result, PRODUCT_ID_TEST)
create_product_method.assert_called_once_with(
parent=parent,
product=product,
product_id=PRODUCT_ID_TEST,
retry=None,
timeout=None,
metadata=None,
)
CloudVisionProductUpdateOperator,
CloudVisionProductDeleteOperator,
)
default_args = {'start_date': airflow.utils.dates.days_ago(1)}
# [START howto_operator_vision_args_common]
GCP_VISION_LOCATION = os.environ.get('GCP_VISION_LOCATION', 'europe-west1')
# [END howto_operator_vision_args_common]
# [START howto_operator_vision_productset]
product_set = ProductSet(display_name='My Product Set 1')
# [END howto_operator_vision_productset]
# [START howto_operator_vision_product]
product = Product(display_name='My Product 1', product_category='toys')
# [END howto_operator_vision_product]
# [START howto_operator_vision_productset_explicit_id]
GCP_VISION_PRODUCT_SET_ID = os.environ.get('GCP_VISION_PRODUCT_SET_ID',
'product_set_explicit_id')
# [END howto_operator_vision_productset_explicit_id]
# [START howto_operator_vision_product_explicit_id]
GCP_VISION_PRODUCT_ID = os.environ.get('GCP_VISION_PRODUCT_ID',
'product_explicit_id')
# [END howto_operator_vision_product_explicit_id]
with models.DAG(
'example_gcp_vision',
default_args=default_args,
CloudVisionProductUpdateOperator,
CloudVisionProductDeleteOperator,
)
try:
# noinspection PyProtectedMember
from unittest import mock
except ImportError:
try:
import mock
except ImportError:
mock = None
PRODUCTSET_TEST = ProductSet(display_name='Test Product Set')
PRODUCTSET_ID_TEST = 'my-productset'
PRODUCT_TEST = Product(display_name='My Product 1', product_category='toys')
PRODUCT_ID_TEST = 'my-product'
LOCATION_TEST = 'europe-west1'
GCP_CONN_ID = 'google_cloud_default'
class CloudVisionProductSetCreateTest(unittest.TestCase):
@mock.patch('airflow.contrib.operators.gcp_vision_operator.CloudVisionHook'
)
def test_minimal_green_path(self, mock_hook):
mock_hook.return_value.create_product_set.return_value = {}
op = CloudVisionProductSetCreateOperator(location=LOCATION_TEST,
product_set=PRODUCTSET_TEST,
task_id='id')
op.execute(context=None)
mock_hook.assert_called_once_with(gcp_conn_id=GCP_CONN_ID)