def setUp(self):
with mock.patch(
'airflow.providers.google.common.hooks.base_google.GoogleBaseHook.__init__',
new=mock_base_gcp_hook_no_default_project_id,
):
self.gcf_function_hook_no_project_id = CloudFunctionsHook(
gcp_conn_id='test', api_version='v1')
def execute(self, context):
hook = CloudFunctionsHook(gcp_conn_id=self.gcp_conn_id, api_version=self.api_version)
try:
return hook.delete_function(self.name)
except HttpError as e:
status = e.resp.status
if status == 404:
self.log.info('The function does not exist in this project')
else:
self.log.error('An error occurred. Exiting.')
raise e
def execute(self, context: Dict):
hook = CloudFunctionsHook(api_version=self.api_version, gcp_conn_id=self.gcp_conn_id)
self.log.info('Calling function %s.', self.function_id)
result = hook.call_function(
function_id=self.function_id,
input_data=self.input_data,
location=self.location,
project_id=self.project_id
)
self.log.info('Function called successfully. Execution id %s', result.get('executionId', None))
self.xcom_push(context=context, key='execution_id', value=result.get('executionId', None))
return result
def execute(self, context):
hook = CloudFunctionsHook(gcp_conn_id=self.gcp_conn_id, api_version=self.api_version)
if self.zip_path_preprocessor.should_upload_function():
self.body[GCF_SOURCE_UPLOAD_URL] = self._upload_source_code(hook)
self._validate_all_body_fields()
self._set_airflow_version_label()
if not self._check_if_function_exists(hook):
self._create_new_function(hook)
else:
self._update_function(hook)
class TestFunctionHookNoDefaultProjectId(unittest.TestCase):
def setUp(self):
with mock.patch(
'airflow.providers.google.cloud.hooks.base.CloudBaseHook.__init__',
new=mock_base_gcp_hook_no_default_project_id):
self.gcf_function_hook_no_project_id = CloudFunctionsHook(
gcp_conn_id='test', api_version='v1')
@mock.patch(
"airflow.providers.google.cloud.hooks.functions.CloudFunctionsHook._authorize"
)
@mock.patch("airflow.providers.google.cloud.hooks.functions.build")
def test_gcf_client_creation(self, mock_build, mock_authorize):
result = self.gcf_function_hook_no_project_id.get_conn()
mock_build.assert_called_once_with('cloudfunctions',
'v1',
http=mock_authorize.return_value,
cache_discovery=False)
self.assertEqual(mock_build.return_value, result)
self.assertEqual(self.gcf_function_hook_no_project_id._conn, result)
@mock.patch(
'airflow.providers.google.cloud.hooks.base.CloudBaseHook.project_id',
new_callable=PropertyMock,
return_value=None)
@mock.patch(
'airflow.providers.google.cloud.hooks.functions.CloudFunctionsHook.get_conn'
)
@mock.patch(
'airflow.providers.google.cloud.hooks.functions.CloudFunctionsHook._wait_for_operation_to_complete'
)
def test_create_new_function_missing_project_id(
self, wait_for_operation_to_complete, get_conn, mock_project_id):
create_method = get_conn.return_value.projects.return_value.locations. \
return_value.functions.return_value.create
execute_method = create_method.return_value.execute
execute_method.return_value = {"name": "operation_id"}
wait_for_operation_to_complete.return_value = None
with self.assertRaises(AirflowException) as cm:
self.gcf_function_hook_no_project_id.create_new_function(
location=GCF_LOCATION, body={})
create_method.assert_not_called()
execute_method.assert_not_called()
err = cm.exception
self.assertIn("The project id must be passed", str(err))
wait_for_operation_to_complete.assert_not_called()
@mock.patch(
'airflow.providers.google.cloud.hooks.functions.CloudFunctionsHook.get_conn'
)
@mock.patch(
'airflow.providers.google.cloud.hooks.functions.CloudFunctionsHook._wait_for_operation_to_complete'
)
def test_create_new_function_overridden_project_id(
self, wait_for_operation_to_complete, get_conn):
create_method = get_conn.return_value.projects.return_value.locations. \
return_value.functions.return_value.create
execute_method = create_method.return_value.execute
execute_method.return_value = {"name": "operation_id"}
wait_for_operation_to_complete.return_value = None
res = self.gcf_function_hook_no_project_id.create_new_function(
project_id=GCP_PROJECT_ID_HOOK_UNIT_TEST,
location=GCF_LOCATION,
body={})
self.assertIsNone(res)
create_method.assert_called_once_with(
body={}, location='projects/example-project/locations/location')
execute_method.assert_called_once_with(num_retries=5)
wait_for_operation_to_complete.assert_called_once_with(
operation_name='operation_id')
@mock.patch(
'airflow.providers.google.cloud.hooks.base.CloudBaseHook.project_id',
new_callable=PropertyMock,
return_value=None)
@mock.patch('requests.put')
@mock.patch(
'airflow.providers.google.cloud.hooks.functions.CloudFunctionsHook.get_conn'
)
def test_upload_function_zip_missing_project_id(self, get_conn,
requests_put,
mock_project_id):
mck = mock.mock_open()
with mock.patch('builtins.open', mck):
generate_upload_url_method = get_conn.return_value.projects.return_value.locations. \
return_value.functions.return_value.generateUploadUrl
execute_method = generate_upload_url_method.return_value.execute
execute_method.return_value = {"uploadUrl": "http://uploadHere"}
requests_put.return_value = None
with self.assertRaises(AirflowException) as cm:
self.gcf_function_hook_no_project_id.upload_function_zip(
location=GCF_LOCATION, zip_path="/tmp/path.zip")
generate_upload_url_method.assert_not_called()
execute_method.assert_not_called()
mck.assert_not_called()
err = cm.exception
self.assertIn("The project id must be passed", str(err))
@mock.patch('requests.put')
@mock.patch(
'airflow.providers.google.cloud.hooks.functions.CloudFunctionsHook.get_conn'
)
def test_upload_function_zip_overridden_project_id(self, get_conn,
requests_put):
mck, open_module = get_open_mock()
with mock.patch('{}.open'.format(open_module), mck):
generate_upload_url_method = get_conn.return_value.projects.return_value.locations. \
return_value.functions.return_value.generateUploadUrl
execute_method = generate_upload_url_method.return_value.execute
execute_method.return_value = {"uploadUrl": "http://uploadHere"}
requests_put.return_value = None
res = self.gcf_function_hook_no_project_id.upload_function_zip(
project_id=GCP_PROJECT_ID_HOOK_UNIT_TEST,
location=GCF_LOCATION,
zip_path="/tmp/path.zip")
self.assertEqual("http://uploadHere", res)
generate_upload_url_method.assert_called_once_with(
parent='projects/example-project/locations/location')
execute_method.assert_called_once_with(num_retries=5)
requests_put.assert_called_once_with(
data=mock.ANY,
headers={
'Content-type': 'application/zip',
'x-goog-content-length-range': '0,104857600'
},
url='http://uploadHere')
class TestFunctionHookDefaultProjectId(unittest.TestCase):
def setUp(self):
with mock.patch(
'airflow.providers.google.cloud.hooks.base.CloudBaseHook.__init__',
new=mock_base_gcp_hook_default_project_id):
self.gcf_function_hook = CloudFunctionsHook(gcp_conn_id='test',
api_version='v1')
@mock.patch(
"airflow.providers.google.cloud.hooks.functions.CloudFunctionsHook._authorize"
)
@mock.patch("airflow.providers.google.cloud.hooks.functions.build")
def test_gcf_client_creation(self, mock_build, mock_authorize):
result = self.gcf_function_hook.get_conn()
mock_build.assert_called_once_with('cloudfunctions',
'v1',
http=mock_authorize.return_value,
cache_discovery=False)
self.assertEqual(mock_build.return_value, result)
self.assertEqual(self.gcf_function_hook._conn, result)
@mock.patch(
'airflow.providers.google.cloud.hooks.base.CloudBaseHook.project_id',
new_callable=PropertyMock,
return_value=GCP_PROJECT_ID_HOOK_UNIT_TEST)
@mock.patch(
'airflow.providers.google.cloud.hooks.functions.CloudFunctionsHook.get_conn'
)
@mock.patch(
'airflow.providers.google.cloud.hooks.functions.CloudFunctionsHook._wait_for_operation_to_complete'
)
def test_create_new_function(self, wait_for_operation_to_complete,
get_conn, mock_project_id):
create_method = get_conn.return_value.projects.return_value.locations.\
return_value.functions.return_value.create
execute_method = create_method.return_value.execute
execute_method.return_value = {"name": "operation_id"}
wait_for_operation_to_complete.return_value = None
res = self.gcf_function_hook.create_new_function(location=GCF_LOCATION,
body={})
self.assertIsNone(res)
create_method.assert_called_once_with(
body={}, location='projects/example-project/locations/location')
execute_method.assert_called_once_with(num_retries=5)
wait_for_operation_to_complete.assert_called_once_with(
operation_name='operation_id')
@mock.patch(
'airflow.providers.google.cloud.hooks.functions.CloudFunctionsHook.get_conn'
)
@mock.patch(
'airflow.providers.google.cloud.hooks.functions.CloudFunctionsHook._wait_for_operation_to_complete'
)
def test_create_new_function_override_project_id(
self, wait_for_operation_to_complete, get_conn):
create_method = get_conn.return_value.projects.return_value.locations. \
return_value.functions.return_value.create
execute_method = create_method.return_value.execute
execute_method.return_value = {"name": "operation_id"}
wait_for_operation_to_complete.return_value = None
res = self.gcf_function_hook.create_new_function(
project_id='new-project', location=GCF_LOCATION, body={})
self.assertIsNone(res)
create_method.assert_called_once_with(
body={}, location='projects/new-project/locations/location')
execute_method.assert_called_once_with(num_retries=5)
wait_for_operation_to_complete.assert_called_once_with(
operation_name='operation_id')
@mock.patch(
'airflow.providers.google.cloud.hooks.functions.CloudFunctionsHook.get_conn'
)
def test_get_function(self, get_conn):
get_method = get_conn.return_value.projects.return_value.locations. \
return_value.functions.return_value.get
execute_method = get_method.return_value.execute
execute_method.return_value = {"name": "function"}
res = self.gcf_function_hook.get_function(name=GCF_FUNCTION)
self.assertIsNotNone(res)
self.assertEqual('function', res['name'])
get_method.assert_called_once_with(name='function')
execute_method.assert_called_once_with(num_retries=5)
@mock.patch(
'airflow.providers.google.cloud.hooks.functions.CloudFunctionsHook.get_conn'
)
@mock.patch(
'airflow.providers.google.cloud.hooks.functions.CloudFunctionsHook._wait_for_operation_to_complete'
)
def test_delete_function(self, wait_for_operation_to_complete, get_conn):
delete_method = get_conn.return_value.projects.return_value.locations. \
return_value.functions.return_value.delete
execute_method = delete_method.return_value.execute
wait_for_operation_to_complete.return_value = None
execute_method.return_value = {"name": "operation_id"}
res = self.gcf_function_hook.delete_function( # pylint: disable=assignment-from-no-return
name=GCF_FUNCTION)
self.assertIsNone(res)
delete_method.assert_called_once_with(name='function')
execute_method.assert_called_once_with(num_retries=5)
@mock.patch(
'airflow.providers.google.cloud.hooks.functions.CloudFunctionsHook.get_conn'
)
@mock.patch(
'airflow.providers.google.cloud.hooks.functions.CloudFunctionsHook._wait_for_operation_to_complete'
)
def test_update_function(self, wait_for_operation_to_complete, get_conn):
patch_method = get_conn.return_value.projects.return_value.locations. \
return_value.functions.return_value.patch
execute_method = patch_method.return_value.execute
execute_method.return_value = {"name": "operation_id"}
wait_for_operation_to_complete.return_value = None
res = self.gcf_function_hook.update_function( # pylint: disable=assignment-from-no-return
update_mask=['a', 'b', 'c'],
name=GCF_FUNCTION,
body={})
self.assertIsNone(res)
patch_method.assert_called_once_with(body={},
name='function',
updateMask='a,b,c')
execute_method.assert_called_once_with(num_retries=5)
wait_for_operation_to_complete.assert_called_once_with(
operation_name='operation_id')
@mock.patch(
'airflow.providers.google.cloud.hooks.base.CloudBaseHook.project_id',
new_callable=PropertyMock,
return_value=GCP_PROJECT_ID_HOOK_UNIT_TEST)
@mock.patch('requests.put')
@mock.patch(
'airflow.providers.google.cloud.hooks.functions.CloudFunctionsHook.get_conn'
)
def test_upload_function_zip(self, get_conn, requests_put,
mock_project_id):
mck, open_module = get_open_mock()
with mock.patch('{}.open'.format(open_module), mck):
generate_upload_url_method = get_conn.return_value.projects.return_value.locations. \
return_value.functions.return_value.generateUploadUrl
execute_method = generate_upload_url_method.return_value.execute
execute_method.return_value = {"uploadUrl": "http://uploadHere"}
requests_put.return_value = None
res = self.gcf_function_hook.upload_function_zip(
location=GCF_LOCATION, zip_path="/tmp/path.zip")
self.assertEqual("http://uploadHere", res)
generate_upload_url_method.assert_called_once_with(
parent='projects/example-project/locations/location')
execute_method.assert_called_once_with(num_retries=5)
requests_put.assert_called_once_with(
data=mock.ANY,
headers={
'Content-type': 'application/zip',
'x-goog-content-length-range': '0,104857600'
},
url='http://uploadHere')
@mock.patch('requests.put')
@mock.patch(
'airflow.providers.google.cloud.hooks.functions.CloudFunctionsHook.get_conn'
)
def test_upload_function_zip_overridden_project_id(self, get_conn,
requests_put):
mck, open_module = get_open_mock()
with mock.patch('{}.open'.format(open_module), mck):
generate_upload_url_method = get_conn.return_value.projects.return_value.locations. \
return_value.functions.return_value.generateUploadUrl
execute_method = generate_upload_url_method.return_value.execute
execute_method.return_value = {"uploadUrl": "http://uploadHere"}
requests_put.return_value = None
res = self.gcf_function_hook.upload_function_zip(
project_id='new-project',
location=GCF_LOCATION,
zip_path="/tmp/path.zip")
self.assertEqual("http://uploadHere", res)
generate_upload_url_method.assert_called_once_with(
parent='projects/new-project/locations/location')
execute_method.assert_called_once_with(num_retries=5)
requests_put.assert_called_once_with(
data=mock.ANY,
headers={
'Content-type': 'application/zip',
'x-goog-content-length-range': '0,104857600'
},
url='http://uploadHere')
@mock.patch(
'airflow.providers.google.cloud.hooks.functions.CloudFunctionsHook.get_conn'
)
def test_call_function(self, mock_get_conn):
payload = {'executionId': 'wh41ppcyoa6l', 'result': 'Hello World!'}
call = mock_get_conn.return_value.projects.return_value.\
locations.return_value.functions.return_value.call
call.return_value.execute.return_value = payload
function_id = "function1234"
input_data = {'key': 'value'}
name = "projects/{project_id}/locations/{location}/functions/{function_id}".format(
project_id=GCP_PROJECT_ID_HOOK_UNIT_TEST,
location=GCF_LOCATION,
function_id=function_id)
result = self.gcf_function_hook.call_function(
function_id=function_id,
location=GCF_LOCATION,
input_data=input_data,
project_id=GCP_PROJECT_ID_HOOK_UNIT_TEST)
call.assert_called_once_with(body=input_data, name=name)
self.assertDictEqual(result, payload)
@mock.patch(
'airflow.providers.google.cloud.hooks.functions.CloudFunctionsHook.get_conn'
)
def test_call_function_error(self, mock_get_conn):
payload = {'error': 'Something very bad'}
call = mock_get_conn.return_value.projects.return_value. \
locations.return_value.functions.return_value.call
call.return_value.execute.return_value = payload
function_id = "function1234"
input_data = {'key': 'value'}
with self.assertRaises(AirflowException):
self.gcf_function_hook.call_function(
function_id=function_id,
location=GCF_LOCATION,
input_data=input_data,
project_id=GCP_PROJECT_ID_HOOK_UNIT_TEST)
class TestFunctionHookNoDefaultProjectId(unittest.TestCase):
def setUp(self):
with mock.patch(
'airflow.providers.google.common.hooks.base_google.GoogleBaseHook.__init__',
new=mock_base_gcp_hook_no_default_project_id,
):
self.gcf_function_hook_no_project_id = CloudFunctionsHook(gcp_conn_id='test', api_version='v1')
@mock.patch("airflow.providers.google.cloud.hooks.functions.CloudFunctionsHook._authorize")
@mock.patch("airflow.providers.google.cloud.hooks.functions.build")
def test_gcf_client_creation(self, mock_build, mock_authorize):
result = self.gcf_function_hook_no_project_id.get_conn()
mock_build.assert_called_once_with(
'cloudfunctions', 'v1', http=mock_authorize.return_value, cache_discovery=False
)
assert mock_build.return_value == result
assert self.gcf_function_hook_no_project_id._conn == result
@mock.patch('airflow.providers.google.cloud.hooks.functions.CloudFunctionsHook.get_conn')
@mock.patch(
'airflow.providers.google.cloud.hooks.functions.CloudFunctionsHook._wait_for_operation_to_complete'
)
def test_create_new_function_overridden_project_id(self, wait_for_operation_to_complete, get_conn):
create_method = (
get_conn.return_value.projects.return_value.locations.return_value.functions.return_value.create
)
execute_method = create_method.return_value.execute
execute_method.return_value = {"name": "operation_id"}
wait_for_operation_to_complete.return_value = None
res = self.gcf_function_hook_no_project_id.create_new_function(
project_id=GCP_PROJECT_ID_HOOK_UNIT_TEST, location=GCF_LOCATION, body={}
)
assert res is None
create_method.assert_called_once_with(body={}, location='projects/example-project/locations/location')
execute_method.assert_called_once_with(num_retries=5)
wait_for_operation_to_complete.assert_called_once_with(operation_name='operation_id')
@mock.patch('requests.put')
@mock.patch('airflow.providers.google.cloud.hooks.functions.CloudFunctionsHook.get_conn')
def test_upload_function_zip_overridden_project_id(self, get_conn, requests_put):
mck, open_module = get_open_mock()
with mock.patch(f'{open_module}.open', mck):
# fmt: off
generate_upload_url_method = get_conn.return_value.projects.return_value.locations. \
return_value.functions.return_value.generateUploadUrl
# fmt: on
execute_method = generate_upload_url_method.return_value.execute
execute_method.return_value = {"uploadUrl": "http://uploadHere"}
requests_put.return_value = None
res = self.gcf_function_hook_no_project_id.upload_function_zip(
project_id=GCP_PROJECT_ID_HOOK_UNIT_TEST, location=GCF_LOCATION, zip_path="/tmp/path.zip"
)
assert "http://uploadHere" == res
generate_upload_url_method.assert_called_once_with(
parent='projects/example-project/locations/location'
)
execute_method.assert_called_once_with(num_retries=5)
requests_put.assert_called_once_with(
data=mock.ANY,
headers={'Content-type': 'application/zip', 'x-goog-content-length-range': '0,104857600'},
url='http://uploadHere',
)
