def execute(self, context: Dict):
hook = CloudMemorystoreHook(gcp_conn_id=self.gcp_conn_id)
result = hook.get_instance(
location=self.location,
instance=self.instance,
project_id=self.project_id,
retry=self.retry,
timeout=self.timeout,
metadata=self.metadata,
)
return MessageToDict(result)
class TestCloudMemorystoreWithDefaultProjectIdHook(TestCase):
def setUp(self,):
with mock.patch(
"airflow.gcp.hooks.cloud_memorystore.CloudMemorystoreHook.__init__",
new=mock_base_gcp_hook_default_project_id,
):
self.hook = CloudMemorystoreHook(gcp_conn_id="test")
@mock.patch(
'airflow.gcp.hooks.base.GoogleCloudBaseHook.project_id',
new_callable=PropertyMock,
return_value=GCP_PROJECT_ID_HOOK_UNIT_TEST
)
@mock.patch( # type: ignore
"airflow.gcp.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(
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.gcp.hooks.base.GoogleCloudBaseHook.project_id',
new_callable=PropertyMock,
return_value=GCP_PROJECT_ID_HOOK_UNIT_TEST
)
@mock.patch( # type: ignore
"airflow.gcp.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(
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.gcp.hooks.base.GoogleCloudBaseHook.project_id',
new_callable=PropertyMock,
return_value=GCP_PROJECT_ID_HOOK_UNIT_TEST
)
@mock.patch( # type: ignore
"airflow.gcp.hooks.cloud_memorystore.CloudMemorystoreHook.get_conn"
)
def test_delete_instance(self, mock_get_conn, mock_project_id):
self.hook.delete_instance(
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.gcp.hooks.base.GoogleCloudBaseHook.project_id',
new_callable=PropertyMock,
return_value=GCP_PROJECT_ID_HOOK_UNIT_TEST
)
@mock.patch( # type: ignore
"airflow.gcp.hooks.cloud_memorystore.CloudMemorystoreHook.get_conn"
)
def test_get_instance(self, mock_get_conn, mock_project_id):
self.hook.get_instance(
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.gcp.hooks.base.GoogleCloudBaseHook.project_id',
new_callable=PropertyMock,
return_value=GCP_PROJECT_ID_HOOK_UNIT_TEST
)
@mock.patch( # type: ignore
"airflow.gcp.hooks.cloud_memorystore.CloudMemorystoreHook.get_conn"
)
def test_list_instances(self, mock_get_conn, mock_project_id):
self.hook.list_instances(
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.gcp.hooks.base.GoogleCloudBaseHook.project_id',
new_callable=PropertyMock,
return_value=GCP_PROJECT_ID_HOOK_UNIT_TEST
)
@mock.patch( # type: ignore
"airflow.gcp.hooks.cloud_memorystore.CloudMemorystoreHook.get_conn"
)
def test_update_instance(self, mock_get_conn, mock_project_id):
self.hook.update_instance(
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,
)
