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):
     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)
 def execute(self, context: Dict):
     hook = CloudMemorystoreHook(gcp_conn_id=self.gcp_conn_id)
     result = hook.list_instances(
         location=self.location,
         page_size=self.page_size,
         project_id=self.project_id,
         retry=self.retry,
         timeout=self.timeout,
         metadata=self.metadata,
     )
     instances = [MessageToDict(a) for a in result]
     return instances
    def execute(self, context: Dict):
        hook = CloudMemorystoreHook(gcp_conn_id=self.gcp_conn_id)

        hook.update_instance(
            update_mask={"paths": ["memory_size_gb"]},
            instance={"memory_size_gb": self.memory_size_gb},
            location=self.location,
            instance_id=self.instance_id,
            project_id=self.project_id,
            retry=self.retry,
            timeout=self.timeout,
            metadata=self.metadata,
        )
 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")
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,
        )