Python BigtableHook.create_instance方法代码示例

# 需要导入模块: from airflow.contrib.hooks.gcp_bigtable_hook import BigtableHook [as 别名]
# 或者: from airflow.contrib.hooks.gcp_bigtable_hook.BigtableHook import create_instance [as 别名]
class TestBigtableHookNoDefaultProjectId(unittest.TestCase):

    def setUp(self):
        with mock.patch('airflow.contrib.hooks.gcp_api_base_hook.GoogleCloudBaseHook.__init__',
                        new=mock_base_gcp_hook_no_default_project_id):
            self.bigtable_hook_no_default_project_id = BigtableHook(gcp_conn_id='test')

    @mock.patch('airflow.contrib.hooks.gcp_bigtable_hook.BigtableHook._get_client')
    def test_get_instance_missing_project_id(self, get_client):
        instance_method = get_client.return_value.instance
        instance_exists_method = instance_method.return_value.exists
        instance_exists_method.return_value = True
        with self.assertRaises(AirflowException) as cm:
            self.bigtable_hook_no_default_project_id.get_instance(instance_id=CBT_INSTANCE)
        instance_exists_method.assert_not_called()
        instance_method.assert_not_called()
        err = cm.exception
        self.assertIn("The project id must be passed", str(err))

    @mock.patch('airflow.contrib.hooks.gcp_bigtable_hook.BigtableHook._get_client')
    def test_get_instance_overridden_project_id(self, get_client):
        instance_method = get_client.return_value.instance
        instance_exists_method = instance_method.return_value.exists
        instance_exists_method.return_value = True
        res = self.bigtable_hook_no_default_project_id.get_instance(
            project_id=GCP_PROJECT_ID_HOOK_UNIT_TEST,
            instance_id=CBT_INSTANCE)
        instance_method.assert_called_once_with('instance')
        instance_exists_method.assert_called_once_with()
        get_client.assert_called_once_with(project_id='example-project')
        self.assertIsNotNone(res)

    @mock.patch('airflow.contrib.hooks.gcp_bigtable_hook.BigtableHook._get_client')
    def test_delete_instance_missing_project_id(self, get_client):
        instance_method = get_client.return_value.instance
        instance_exists_method = instance_method.return_value.exists
        delete_method = instance_method.return_value.delete
        instance_exists_method.return_value = True
        with self.assertRaises(AirflowException) as cm:
            self.bigtable_hook_no_default_project_id.delete_instance(instance_id=CBT_INSTANCE)
        instance_exists_method.assert_not_called()
        instance_method.assert_not_called()
        delete_method.assert_not_called()
        err = cm.exception
        self.assertIn("The project id must be passed", str(err))

    @mock.patch('airflow.contrib.hooks.gcp_bigtable_hook.BigtableHook._get_client')
    def test_delete_instance_overridden_project_id(self, get_client):
        instance_method = get_client.return_value.instance
        instance_exists_method = instance_method.return_value.exists
        instance_exists_method.return_value = True
        delete_method = instance_method.return_value.delete
        res = self.bigtable_hook_no_default_project_id.delete_instance(
            project_id=GCP_PROJECT_ID_HOOK_UNIT_TEST, instance_id=CBT_INSTANCE)
        instance_method.assert_called_once_with('instance')
        instance_exists_method.assert_called_once_with()
        delete_method.assert_called_once_with()
        get_client.assert_called_once_with(project_id='example-project')
        self.assertIsNone(res)

    @mock.patch('google.cloud.bigtable.instance.Instance.create')
    @mock.patch('airflow.contrib.hooks.gcp_bigtable_hook.BigtableHook._get_client')
    def test_create_instance_missing_project_id(self, get_client, instance_create):
        operation = mock.Mock()
        operation.result_return_value = Instance(instance_id=CBT_INSTANCE, client=get_client)
        instance_create.return_value = operation
        with self.assertRaises(AirflowException) as cm:
            self.bigtable_hook_no_default_project_id.create_instance(
                instance_id=CBT_INSTANCE,
                main_cluster_id=CBT_CLUSTER,
                main_cluster_zone=CBT_ZONE)
        get_client.assert_not_called()
        instance_create.assert_not_called()
        err = cm.exception
        self.assertIn("The project id must be passed", str(err))

    @mock.patch('google.cloud.bigtable.instance.Instance.create')
    @mock.patch('airflow.contrib.hooks.gcp_bigtable_hook.BigtableHook._get_client')
    def test_create_instance_overridden_project_id(self, get_client, instance_create):
        operation = mock.Mock()
        operation.result_return_value = Instance(instance_id=CBT_INSTANCE, client=get_client)
        instance_create.return_value = operation
        res = self.bigtable_hook_no_default_project_id.create_instance(
            project_id=GCP_PROJECT_ID_HOOK_UNIT_TEST,
            instance_id=CBT_INSTANCE,
            main_cluster_id=CBT_CLUSTER,
            main_cluster_zone=CBT_ZONE)
        get_client.assert_called_once_with(project_id='example-project')
        instance_create.assert_called_once_with(clusters=mock.ANY)
        self.assertEqual(res.instance_id, 'instance')

    @mock.patch('airflow.contrib.hooks.gcp_bigtable_hook.BigtableHook._get_client')
    def test_delete_table_missing_project_id(self, get_client):
        instance_method = get_client.return_value.instance
        instance_exists_method = instance_method.return_value.exists
        table_delete_method = instance_method.return_value.table.return_value.delete
        instance_exists_method.return_value = True
        with self.assertRaises(AirflowException) as cm:
            self.bigtable_hook_no_default_project_id.delete_table(
                instance_id=CBT_INSTANCE,
#.........这里部分代码省略.........

# 需要导入模块: from airflow.contrib.hooks.gcp_bigtable_hook import BigtableHook [as 别名]
# 或者: from airflow.contrib.hooks.gcp_bigtable_hook.BigtableHook import create_instance [as 别名]
class BigtableInstanceCreateOperator(BaseOperator, BigtableValidationMixin):
    """
    Creates a new Cloud Bigtable instance.
    If the Cloud Bigtable instance with the given ID exists, the operator does not
    compare its configuration
    and immediately succeeds. No changes are made to the existing instance.

    For more details about instance creation have a look at the reference:
    https://googleapis.github.io/google-cloud-python/latest/bigtable/instance.html#google.cloud.bigtable.instance.Instance.create

    :type instance_id: str
    :param instance_id: The ID of the Cloud Bigtable instance to create.
    :type main_cluster_id: str
    :param main_cluster_id: The ID for main cluster for the new instance.
    :type main_cluster_zone: str
    :param main_cluster_zone: The zone for main cluster
        See https://cloud.google.com/bigtable/docs/locations for more details.
    :type project_id: str
    :param project_id: Optional, the ID of the GCP project.  If set to None or missing,
            the default project_id from the GCP connection is used.
    :type replica_cluster_id: str
    :param replica_cluster_id: (optional) The ID for replica cluster for the new instance.
    :type replica_cluster_zone: str
    :param replica_cluster_zone: (optional)  The zone for replica cluster.
    :type instance_type: IntEnum
    :param instance_type: (optional) The type of the instance.
    :type instance_display_name: str
    :param instance_display_name: (optional) Human-readable name of the instance. Defaults
        to ``instance_id``.
    :type instance_labels: dict
    :param instance_labels: (optional) Dictionary of labels to associate
        with the instance.
    :type cluster_nodes: int
    :param cluster_nodes: (optional) Number of nodes for cluster.
    :type cluster_storage_type: IntEnum
    :param cluster_storage_type: (optional) The type of storage.
    :type timeout: int
    :param timeout: (optional) timeout (in seconds) for instance creation.
                    If None is not specified, Operator will wait indefinitely.
    """

    REQUIRED_ATTRIBUTES = ('instance_id', 'main_cluster_id',
                           'main_cluster_zone')
    template_fields = ['project_id', 'instance_id', 'main_cluster_id',
                       'main_cluster_zone']

    @apply_defaults
    def __init__(self,
                 instance_id,
                 main_cluster_id,
                 main_cluster_zone,
                 project_id=None,
                 replica_cluster_id=None,
                 replica_cluster_zone=None,
                 instance_display_name=None,
                 instance_type=None,
                 instance_labels=None,
                 cluster_nodes=None,
                 cluster_storage_type=None,
                 timeout=None,
                 *args, **kwargs):
        self.project_id = project_id
        self.instance_id = instance_id
        self.main_cluster_id = main_cluster_id
        self.main_cluster_zone = main_cluster_zone
        self.replica_cluster_id = replica_cluster_id
        self.replica_cluster_zone = replica_cluster_zone
        self.instance_display_name = instance_display_name
        self.instance_type = instance_type
        self.instance_labels = instance_labels
        self.cluster_nodes = cluster_nodes
        self.cluster_storage_type = cluster_storage_type
        self.timeout = timeout
        self._validate_inputs()
        self.hook = BigtableHook()
        super(BigtableInstanceCreateOperator, self).__init__(*args, **kwargs)

    def execute(self, context):
        instance = self.hook.get_instance(project_id=self.project_id,
                                          instance_id=self.instance_id)
        if instance:
            # Based on Instance.__eq__ instance with the same ID and client is
            # considered as equal.
            self.log.info(
                "The instance '%s' already exists in this project. "
                "Consider it as created",
                self.instance_id
            )
            return
        try:
            self.hook.create_instance(
                project_id=self.project_id,
                instance_id=self.instance_id,
                main_cluster_id=self.main_cluster_id,
                main_cluster_zone=self.main_cluster_zone,
                replica_cluster_id=self.replica_cluster_id,
                replica_cluster_zone=self.replica_cluster_zone,
                instance_display_name=self.instance_display_name,
                instance_type=self.instance_type,
                instance_labels=self.instance_labels,
#.........这里部分代码省略.........

# 需要导入模块: from airflow.contrib.hooks.gcp_bigtable_hook import BigtableHook [as 别名]
# 或者: from airflow.contrib.hooks.gcp_bigtable_hook.BigtableHook import create_instance [as 别名]
class TestBigtableHookDefaultProjectId(unittest.TestCase):

    def setUp(self):
        with mock.patch('airflow.contrib.hooks.gcp_api_base_hook.GoogleCloudBaseHook.__init__',
                        new=mock_base_gcp_hook_default_project_id):
            self.bigtable_hook_default_project_id = BigtableHook(gcp_conn_id='test')

    @mock.patch('airflow.contrib.hooks.gcp_bigtable_hook.BigtableHook._get_client')
    def test_get_instance(self, get_client):
        instance_method = get_client.return_value.instance
        instance_exists_method = instance_method.return_value.exists
        instance_exists_method.return_value = True
        res = self.bigtable_hook_default_project_id.get_instance(
            instance_id=CBT_INSTANCE)
        instance_method.assert_called_once_with('instance')
        instance_exists_method.assert_called_once_with()
        get_client.assert_called_once_with(project_id='example-project')
        self.assertIsNotNone(res)

    @mock.patch('airflow.contrib.hooks.gcp_bigtable_hook.BigtableHook._get_client')
    def test_get_instance_overridden_project_id(self, get_client):
        instance_method = get_client.return_value.instance
        instance_exists_method = instance_method.return_value.exists
        instance_exists_method.return_value = True
        res = self.bigtable_hook_default_project_id.get_instance(
            project_id='new-project',
            instance_id=CBT_INSTANCE)
        instance_method.assert_called_once_with('instance')
        instance_exists_method.assert_called_once_with()
        get_client.assert_called_once_with(project_id='new-project')
        self.assertIsNotNone(res)

    @mock.patch('airflow.contrib.hooks.gcp_bigtable_hook.BigtableHook._get_client')
    def test_get_instance_no_instance(self, get_client):
        instance_method = get_client.return_value.instance
        instance_exists_method = instance_method.return_value.exists
        instance_exists_method.return_value = False
        res = self.bigtable_hook_default_project_id.get_instance(
            instance_id=CBT_INSTANCE)
        instance_method.assert_called_once_with('instance')
        instance_exists_method.assert_called_once_with()
        get_client.assert_called_once_with(project_id='example-project')
        self.assertIsNone(res)

    @mock.patch('airflow.contrib.hooks.gcp_bigtable_hook.BigtableHook._get_client')
    def test_delete_instance(self, get_client):
        instance_method = get_client.return_value.instance
        instance_exists_method = instance_method.return_value.exists
        instance_exists_method.return_value = True
        delete_method = instance_method.return_value.delete
        res = self.bigtable_hook_default_project_id.delete_instance(
            instance_id=CBT_INSTANCE)
        instance_method.assert_called_once_with('instance')
        instance_exists_method.assert_called_once_with()
        delete_method.assert_called_once_with()
        get_client.assert_called_once_with(project_id='example-project')
        self.assertIsNone(res)

    @mock.patch('airflow.contrib.hooks.gcp_bigtable_hook.BigtableHook._get_client')
    def test_delete_instance_overridden_project_id(self, get_client):
        instance_method = get_client.return_value.instance
        instance_exists_method = instance_method.return_value.exists
        instance_exists_method.return_value = True
        delete_method = instance_method.return_value.delete
        res = self.bigtable_hook_default_project_id.delete_instance(
            project_id='new-project', instance_id=CBT_INSTANCE)
        instance_method.assert_called_once_with('instance')
        instance_exists_method.assert_called_once_with()
        delete_method.assert_called_once_with()
        get_client.assert_called_once_with(project_id='new-project')
        self.assertIsNone(res)

    @mock.patch('airflow.contrib.hooks.gcp_bigtable_hook.BigtableHook._get_client')
    def test_delete_instance_no_instance(self, get_client):
        instance_method = get_client.return_value.instance
        instance_exists_method = instance_method.return_value.exists
        instance_exists_method.return_value = False
        delete_method = instance_method.return_value.delete
        self.bigtable_hook_default_project_id.delete_instance(
            instance_id=CBT_INSTANCE)
        instance_method.assert_called_once_with('instance')
        instance_exists_method.assert_called_once_with()
        delete_method.assert_not_called()
        get_client.assert_called_once_with(project_id='example-project')

    @mock.patch('google.cloud.bigtable.instance.Instance.create')
    @mock.patch('airflow.contrib.hooks.gcp_bigtable_hook.BigtableHook._get_client')
    def test_create_instance(self, get_client, instance_create):
        operation = mock.Mock()
        operation.result_return_value = Instance(instance_id=CBT_INSTANCE, client=get_client)
        instance_create.return_value = operation
        res = self.bigtable_hook_default_project_id.create_instance(
            instance_id=CBT_INSTANCE,
            main_cluster_id=CBT_CLUSTER,
            main_cluster_zone=CBT_ZONE)
        get_client.assert_called_once_with(project_id='example-project')
        instance_create.assert_called_once_with(clusters=mock.ANY)
        self.assertEqual(res.instance_id, 'instance')

    @mock.patch('google.cloud.bigtable.instance.Instance.create')
#.........这里部分代码省略.........