def execute(self, context: Dict):
hook = CloudMemorystoreHook(gcp_conn_id=self.gcp_conn_id)
hook.delete_instance(
location=self.location,
instance=self.instance,
project_id=self.project_id,
retry=self.retry,
timeout=self.timeout,
metadata=self.metadata,
)
def execute(self, context: dict) -> None:
hook = CloudMemorystoreHook(
gcp_conn_id=self.gcp_conn_id, impersonation_chain=self.impersonation_chain
)
hook.delete_instance(
location=self.location,
instance=self.instance,
project_id=self.project_id,
retry=self.retry,
timeout=self.timeout,
metadata=self.metadata,
)
class TestCloudMemorystoreWithDefaultProjectIdHook(TestCase):
def setUp(self, ):
with mock.patch(
"airflow.providers.google.cloud.hooks.cloud_memorystore.CloudMemorystoreHook.__init__",
new=mock_base_gcp_hook_default_project_id,
):
self.hook = CloudMemorystoreHook(gcp_conn_id="test")
@mock.patch(
'airflow.providers.google.common.hooks.base_google.GoogleBaseHook.project_id',
new_callable=PropertyMock,
return_value=GCP_PROJECT_ID_HOOK_UNIT_TEST)
@mock.patch( # type: ignore
"airflow.providers.google.cloud.hooks.cloud_memorystore.CloudMemorystoreHook.get_conn"
)
def test_create_instance_when_exists(self, mock_get_conn, mock_project_id):
mock_get_conn.return_value.get_instance.return_value = Instance(
name=TEST_NAME)
result = self.hook.create_instance( # pylint: disable=no-value-for-parameter
location=TEST_LOCATION,
instance_id=TEST_INSTANCE_ID,
instance=Instance(name=TEST_NAME),
retry=TEST_RETRY,
timeout=TEST_TIMEOUT,
metadata=TEST_METADATA,
)
mock_get_conn.return_value.get_instance.assert_called_once_with(
name=TEST_NAME_DEFAULT_PROJECT_ID,
retry=TEST_RETRY,
timeout=TEST_TIMEOUT,
metadata=TEST_METADATA)
self.assertEqual(Instance(name=TEST_NAME), result)
@mock.patch(
'airflow.providers.google.common.hooks.base_google.GoogleBaseHook.project_id',
new_callable=PropertyMock,
return_value=GCP_PROJECT_ID_HOOK_UNIT_TEST)
@mock.patch( # type: ignore
"airflow.providers.google.cloud.hooks.cloud_memorystore.CloudMemorystoreHook.get_conn"
)
def test_create_instance_when_not_exists(self, mock_get_conn,
mock_project_id):
mock_get_conn.return_value.get_instance.side_effect = [
NotFound("Instnace not found"),
Instance(name=TEST_NAME),
]
mock_get_conn.return_value.create_instance.return_value.result.return_value = Instance(
name=TEST_NAME)
result = self.hook.create_instance( # pylint: disable=no-value-for-parameter
location=TEST_LOCATION,
instance_id=TEST_INSTANCE_ID,
instance=Instance(name=TEST_NAME),
retry=TEST_RETRY,
timeout=TEST_TIMEOUT,
metadata=TEST_METADATA,
)
mock_get_conn.return_value.get_instance.has_calls([
mock.call(name=TEST_NAME,
retry=TEST_RETRY,
timeout=TEST_TIMEOUT,
metadata=TEST_METADATA),
mock.call(name=TEST_NAME,
retry=TEST_RETRY,
timeout=TEST_TIMEOUT,
metadata=TEST_METADATA),
])
mock_get_conn.return_value.create_instance.assert_called_once_with(
instance=Instance(
name=TEST_NAME,
labels={
"airflow-version":
"v" + version.version.replace(".", "-").replace("+", "-")
},
),
instance_id=TEST_INSTANCE_ID,
metadata=TEST_METADATA,
parent=TEST_PARENT_DEFAULT_PROJECT_ID,
retry=TEST_RETRY,
timeout=TEST_TIMEOUT,
)
self.assertEqual(Instance(name=TEST_NAME), result)
@mock.patch(
'airflow.providers.google.common.hooks.base_google.GoogleBaseHook.project_id',
new_callable=PropertyMock,
return_value=GCP_PROJECT_ID_HOOK_UNIT_TEST)
@mock.patch( # type: ignore
"airflow.providers.google.cloud.hooks.cloud_memorystore.CloudMemorystoreHook.get_conn"
)
def test_delete_instance(self, mock_get_conn, mock_project_id):
self.hook.delete_instance( # pylint: disable=no-value-for-parameter
location=TEST_LOCATION,
instance=TEST_INSTANCE_ID,
retry=TEST_RETRY,
timeout=TEST_TIMEOUT,
metadata=TEST_METADATA,
)
mock_get_conn.return_value.delete_instance.assert_called_once_with(
name=TEST_NAME_DEFAULT_PROJECT_ID,
retry=TEST_RETRY,
timeout=TEST_TIMEOUT,
metadata=TEST_METADATA)
@mock.patch(
'airflow.providers.google.common.hooks.base_google.GoogleBaseHook.project_id',
new_callable=PropertyMock,
return_value=GCP_PROJECT_ID_HOOK_UNIT_TEST)
@mock.patch( # type: ignore
"airflow.providers.google.cloud.hooks.cloud_memorystore.CloudMemorystoreHook.get_conn"
)
def test_get_instance(self, mock_get_conn, mock_project_id):
self.hook.get_instance( # pylint: disable=no-value-for-parameter
location=TEST_LOCATION,
instance=TEST_INSTANCE_ID,
retry=TEST_RETRY,
timeout=TEST_TIMEOUT,
metadata=TEST_METADATA,
)
mock_get_conn.return_value.get_instance.assert_called_once_with(
name=TEST_NAME_DEFAULT_PROJECT_ID,
retry=TEST_RETRY,
timeout=TEST_TIMEOUT,
metadata=TEST_METADATA)
@mock.patch(
'airflow.providers.google.common.hooks.base_google.GoogleBaseHook.project_id',
new_callable=PropertyMock,
return_value=GCP_PROJECT_ID_HOOK_UNIT_TEST)
@mock.patch( # type: ignore
"airflow.providers.google.cloud.hooks.cloud_memorystore.CloudMemorystoreHook.get_conn"
)
def test_list_instances(self, mock_get_conn, mock_project_id):
self.hook.list_instances( # pylint: disable=no-value-for-parameter
location=TEST_LOCATION,
page_size=TEST_PAGE_SIZE,
retry=TEST_RETRY,
timeout=TEST_TIMEOUT,
metadata=TEST_METADATA,
)
mock_get_conn.return_value.list_instances.assert_called_once_with(
parent=TEST_PARENT_DEFAULT_PROJECT_ID,
page_size=TEST_PAGE_SIZE,
retry=TEST_RETRY,
timeout=TEST_TIMEOUT,
metadata=TEST_METADATA,
)
@mock.patch(
'airflow.providers.google.common.hooks.base_google.GoogleBaseHook.project_id',
new_callable=PropertyMock,
return_value=GCP_PROJECT_ID_HOOK_UNIT_TEST)
@mock.patch( # type: ignore
"airflow.providers.google.cloud.hooks.cloud_memorystore.CloudMemorystoreHook.get_conn"
)
def test_update_instance(self, mock_get_conn, mock_project_id):
self.hook.update_instance( # pylint: disable=no-value-for-parameter
update_mask=TEST_UPDATE_MASK,
instance=Instance(name=TEST_NAME),
location=TEST_LOCATION,
instance_id=TEST_INSTANCE_ID,
retry=TEST_RETRY,
timeout=TEST_TIMEOUT,
metadata=TEST_METADATA,
)
mock_get_conn.return_value.update_instance.assert_called_once_with(
update_mask=TEST_UPDATE_MASK,
instance=Instance(name=TEST_NAME_DEFAULT_PROJECT_ID),
retry=TEST_RETRY,
timeout=TEST_TIMEOUT,
metadata=TEST_METADATA,
)
class TestCloudMemorystoreWithoutDefaultProjectIdHook(TestCase):
def setUp(self, ):
with mock.patch(
"airflow.providers.google.cloud.hooks.cloud_memorystore.CloudMemorystoreHook.__init__",
new=mock_base_gcp_hook_no_default_project_id,
):
self.hook = CloudMemorystoreHook(gcp_conn_id="test")
@mock.patch(
"airflow.providers.google.cloud.hooks.cloud_memorystore.CloudMemorystoreHook.get_conn"
)
def test_create_instance_when_exists(self, mock_get_conn):
mock_get_conn.return_value.get_instance.return_value = Instance(
name=TEST_NAME)
result = self.hook.create_instance(
location=TEST_LOCATION,
instance_id=TEST_INSTANCE_ID,
instance=Instance(name=TEST_NAME),
project_id=TEST_PROJECT_ID,
retry=TEST_RETRY,
timeout=TEST_TIMEOUT,
metadata=TEST_METADATA,
)
mock_get_conn.return_value.get_instance.assert_called_once_with(
request=dict(
name=
"projects/test-project-id/locations/test-location/instances/test-instance-id"
),
retry=TEST_RETRY,
timeout=TEST_TIMEOUT,
metadata=TEST_METADATA,
)
assert Instance(name=TEST_NAME) == result
@mock.patch(
"airflow.providers.google.cloud.hooks.cloud_memorystore.CloudMemorystoreHook.get_conn"
)
def test_create_instance_when_not_exists(self, mock_get_conn):
mock_get_conn.return_value.get_instance.side_effect = [
NotFound("Instance not found"),
Instance(name=TEST_NAME),
]
mock_get_conn.return_value.create_instance.return_value.result.return_value = Instance(
name=TEST_NAME)
result = self.hook.create_instance(
location=TEST_LOCATION,
instance_id=TEST_INSTANCE_ID,
instance=Instance(name=TEST_NAME),
project_id=TEST_PROJECT_ID,
retry=TEST_RETRY,
timeout=TEST_TIMEOUT,
metadata=TEST_METADATA,
)
mock_get_conn.return_value.get_instance.has_calls([
mock.call(
name=
"projects/test-project-id/locations/test-location/instances/test-instance-id",
retry=TEST_RETRY,
timeout=TEST_TIMEOUT,
metadata=TEST_METADATA,
),
mock.call(
name=
"projects/test-project-id/locations/test-location/instances/test-instance-id",
retry=TEST_RETRY,
timeout=TEST_TIMEOUT,
metadata=TEST_METADATA,
),
])
mock_get_conn.return_value.create_instance.assert_called_once_with(
request=dict(
parent=TEST_PARENT,
instance=Instance(
name=TEST_NAME,
labels={
"airflow-version":
"v" +
version.version.replace(".", "-").replace("+", "-")
},
),
instance_id=TEST_INSTANCE_ID,
),
metadata=TEST_METADATA,
retry=TEST_RETRY,
timeout=TEST_TIMEOUT,
)
assert Instance(name=TEST_NAME) == result
@mock.patch(
'airflow.providers.google.common.hooks.base_google.GoogleBaseHook.project_id',
new_callable=PropertyMock,
return_value=None,
)
@mock.patch(
"airflow.providers.google.cloud.hooks.cloud_memorystore.CloudMemorystoreHook.get_conn"
)
def test_create_instance_without_project_id(self, mock_get_conn,
mock_project_id):
with pytest.raises(AirflowException):
self.hook.create_instance(
location=TEST_LOCATION,
instance_id=TEST_INSTANCE_ID,
instance=Instance(name=TEST_NAME),
project_id=None,
retry=TEST_RETRY,
timeout=TEST_TIMEOUT,
metadata=TEST_METADATA,
)
@mock.patch(
"airflow.providers.google.cloud.hooks.cloud_memorystore.CloudMemorystoreHook.get_conn"
)
def test_delete_instance(self, mock_get_conn):
self.hook.delete_instance(
location=TEST_LOCATION,
instance=TEST_INSTANCE_ID,
project_id=TEST_PROJECT_ID,
retry=TEST_RETRY,
timeout=TEST_TIMEOUT,
metadata=TEST_METADATA,
)
mock_get_conn.return_value.delete_instance.assert_called_once_with(
request=dict(name=TEST_NAME),
retry=TEST_RETRY,
timeout=TEST_TIMEOUT,
metadata=TEST_METADATA)
@mock.patch(
'airflow.providers.google.common.hooks.base_google.GoogleBaseHook.project_id',
new_callable=PropertyMock,
return_value=None,
)
@mock.patch(
"airflow.providers.google.cloud.hooks.cloud_memorystore.CloudMemorystoreHook.get_conn"
)
def test_delete_instance_without_project_id(self, mock_get_conn,
mock_project_id):
with pytest.raises(AirflowException):
self.hook.delete_instance(
location=TEST_LOCATION,
instance=Instance(name=TEST_NAME),
project_id=None,
retry=TEST_RETRY,
timeout=TEST_TIMEOUT,
metadata=TEST_METADATA,
)
@mock.patch(
"airflow.providers.google.cloud.hooks.cloud_memorystore.CloudMemorystoreHook.get_conn"
)
def test_get_instance(self, mock_get_conn):
self.hook.get_instance(
location=TEST_LOCATION,
instance=TEST_INSTANCE_ID,
project_id=TEST_PROJECT_ID,
retry=TEST_RETRY,
timeout=TEST_TIMEOUT,
metadata=TEST_METADATA,
)
mock_get_conn.return_value.get_instance.assert_called_once_with(
request=dict(name=TEST_NAME),
retry=TEST_RETRY,
timeout=TEST_TIMEOUT,
metadata=TEST_METADATA)
@mock.patch(
'airflow.providers.google.common.hooks.base_google.GoogleBaseHook.project_id',
new_callable=PropertyMock,
return_value=None,
)
@mock.patch(
"airflow.providers.google.cloud.hooks.cloud_memorystore.CloudMemorystoreHook.get_conn"
)
def test_get_instance_without_project_id(self, mock_get_conn,
mock_project_id):
with pytest.raises(AirflowException):
self.hook.get_instance(
location=TEST_LOCATION,
instance=Instance(name=TEST_NAME),
project_id=None,
retry=TEST_RETRY,
timeout=TEST_TIMEOUT,
metadata=TEST_METADATA,
)
@mock.patch(
"airflow.providers.google.cloud.hooks.cloud_memorystore.CloudMemorystoreHook.get_conn"
)
def test_list_instances(self, mock_get_conn):
self.hook.list_instances(
location=TEST_LOCATION,
page_size=TEST_PAGE_SIZE,
project_id=TEST_PROJECT_ID,
retry=TEST_RETRY,
timeout=TEST_TIMEOUT,
metadata=TEST_METADATA,
)
mock_get_conn.return_value.list_instances.assert_called_once_with(
request=dict(parent=TEST_PARENT, page_size=TEST_PAGE_SIZE),
retry=TEST_RETRY,
timeout=TEST_TIMEOUT,
metadata=TEST_METADATA,
)
@mock.patch(
'airflow.providers.google.common.hooks.base_google.GoogleBaseHook.project_id',
new_callable=PropertyMock,
return_value=None,
)
@mock.patch(
"airflow.providers.google.cloud.hooks.cloud_memorystore.CloudMemorystoreHook.get_conn"
)
def test_list_instances_without_project_id(self, mock_get_conn,
mock_project_id):
with pytest.raises(AirflowException):
self.hook.list_instances(
location=TEST_LOCATION,
page_size=TEST_PAGE_SIZE,
project_id=None,
retry=TEST_RETRY,
timeout=TEST_TIMEOUT,
metadata=TEST_METADATA,
)
@mock.patch(
"airflow.providers.google.cloud.hooks.cloud_memorystore.CloudMemorystoreHook.get_conn"
)
def test_update_instance(self, mock_get_conn):
self.hook.update_instance(
update_mask=TEST_UPDATE_MASK,
instance=Instance(name=TEST_NAME),
retry=TEST_RETRY,
timeout=TEST_TIMEOUT,
metadata=TEST_METADATA,
project_id=TEST_PROJECT_ID,
)
mock_get_conn.return_value.update_instance.assert_called_once_with(
request=dict(update_mask={'paths': ['memory_size_gb']},
instance=Instance(name=TEST_NAME)),
retry=TEST_RETRY,
timeout=TEST_TIMEOUT,
metadata=TEST_METADATA,
)
@mock.patch(
'airflow.providers.google.common.hooks.base_google.GoogleBaseHook.project_id',
new_callable=PropertyMock,
return_value=None,
)
@mock.patch(
"airflow.providers.google.cloud.hooks.cloud_memorystore.CloudMemorystoreHook.get_conn"
)
def test_update_instance_without_project_id(self, mock_get_conn,
mock_project_id):
with pytest.raises(AirflowException):
self.hook.update_instance( # pylint: disable=no-value-for-parameter
update_mask=TEST_UPDATE_MASK,
instance=Instance(name=TEST_NAME),
retry=TEST_RETRY,
timeout=TEST_TIMEOUT,
metadata=TEST_METADATA,
)
