def test_list_clusters(self):
from google.cloud.bigtable_admin_v2.services.bigtable_instance_admin import (
BigtableInstanceAdminClient, )
from google.cloud.bigtable_admin_v2.types import (
bigtable_instance_admin as messages_v2_pb2, )
from google.cloud.bigtable_admin_v2.types import instance as data_v2_pb2
from google.cloud.bigtable.instance import Instance
from google.cloud.bigtable.instance import Cluster
credentials = _make_credentials()
client = self._make_client(project=self.PROJECT,
credentials=credentials,
admin=True)
instance = Instance(self.INSTANCE_ID, client)
failed_location = "FAILED"
cluster_id1 = "cluster-id1"
cluster_id2 = "cluster-id2"
cluster_path_template = "projects/{}/instances/{}/clusters/{}"
cluster_name1 = cluster_path_template.format(self.PROJECT,
self.INSTANCE_ID,
cluster_id1)
cluster_name2 = cluster_path_template.format(self.PROJECT,
self.INSTANCE_ID,
cluster_id2)
# Create response_pb
response_pb = messages_v2_pb2.ListClustersResponse(
failed_locations=[failed_location],
clusters=[
data_v2_pb2.Cluster(name=cluster_name1),
data_v2_pb2.Cluster(name=cluster_name2),
],
)
# Patch the stub used by the API method.
instance_api = mock.create_autospec(BigtableInstanceAdminClient)
instance_api.list_clusters.side_effect = [response_pb]
instance_api.cluster_path = cluster_path_template.format
client._instance_admin_client = instance_api
# Perform the method and check the result.
clusters, failed_locations = instance.list_clusters()
cluster_1, cluster_2 = clusters
self.assertIsInstance(cluster_1, Cluster)
self.assertEqual(cluster_1.name, cluster_name1)
self.assertIsInstance(cluster_2, Cluster)
self.assertEqual(cluster_2.name, cluster_name2)
self.assertEqual(failed_locations, [failed_location])
def test_from_pb_success(self):
from google.cloud.bigtable_admin_v2.types import instance as data_v2_pb2
from google.cloud.bigtable import enums
client = _Client(self.PROJECT)
instance = _Instance(self.INSTANCE_ID, client)
location = self.LOCATION_PATH + self.LOCATION_ID
state = enums.Cluster.State.RESIZING
storage_type = enums.StorageType.SSD
cluster_pb = data_v2_pb2.Cluster(
name=self.CLUSTER_NAME,
location=location,
state=state,
serve_nodes=self.SERVE_NODES,
default_storage_type=storage_type,
)
klass = self._get_target_class()
cluster = klass.from_pb(cluster_pb, instance)
self.assertIsInstance(cluster, klass)
self.assertEqual(cluster._instance, instance)
self.assertEqual(cluster.cluster_id, self.CLUSTER_ID)
self.assertEqual(cluster.location_id, self.LOCATION_ID)
self.assertEqual(cluster.state, state)
self.assertEqual(cluster.serve_nodes, self.SERVE_NODES)
self.assertEqual(cluster.default_storage_type, storage_type)
def test_cluster_from_pb_success():
from google.cloud.bigtable_admin_v2.types import instance as data_v2_pb2
from google.cloud.bigtable.cluster import Cluster
from google.cloud.bigtable import enums
client = _Client(PROJECT)
instance = _Instance(INSTANCE_ID, client)
location = LOCATION_PATH + LOCATION_ID
state = enums.Cluster.State.RESIZING
storage_type = enums.StorageType.SSD
cluster_pb = data_v2_pb2.Cluster(
name=CLUSTER_NAME,
location=location,
state=state,
serve_nodes=SERVE_NODES,
default_storage_type=storage_type,
encryption_config=data_v2_pb2.Cluster.EncryptionConfig(
kms_key_name=KMS_KEY_NAME, ),
)
cluster = Cluster.from_pb(cluster_pb, instance)
assert isinstance(cluster, Cluster)
assert cluster._instance == instance
assert cluster.cluster_id == CLUSTER_ID
assert cluster.location_id == LOCATION_ID
assert cluster.state == state
assert cluster.serve_nodes == SERVE_NODES
assert cluster.default_storage_type == storage_type
assert cluster.kms_key_name == KMS_KEY_NAME
def test_reload(self):
from google.cloud.bigtable_admin_v2.services.bigtable_instance_admin import (
BigtableInstanceAdminClient,
)
from google.cloud.bigtable_admin_v2.types import instance as data_v2_pb2
from google.cloud.bigtable.enums import StorageType
from google.cloud.bigtable.enums import Cluster
api = mock.create_autospec(BigtableInstanceAdminClient)
credentials = _make_credentials()
client = self._make_client(
project=self.PROJECT, credentials=credentials, admin=True
)
STORAGE_TYPE_SSD = StorageType.SSD
instance = _Instance(self.INSTANCE_ID, client)
cluster = self._make_one(
self.CLUSTER_ID,
instance,
location_id=self.LOCATION_ID,
serve_nodes=self.SERVE_NODES,
default_storage_type=STORAGE_TYPE_SSD,
)
# Create response_pb
LOCATION_ID_FROM_SERVER = "new-location-id"
STATE = Cluster.State.READY
SERVE_NODES_FROM_SERVER = 10
STORAGE_TYPE_FROM_SERVER = StorageType.HDD
response_pb = data_v2_pb2.Cluster(
name=cluster.name,
location=self.LOCATION_PATH + LOCATION_ID_FROM_SERVER,
state=STATE,
serve_nodes=SERVE_NODES_FROM_SERVER,
default_storage_type=STORAGE_TYPE_FROM_SERVER,
)
# Patch the stub used by the API method.
client._instance_admin_client = api
instance_stub = client._instance_admin_client
instance_stub.get_cluster.side_effect = [response_pb]
# Create expected_result.
expected_result = None # reload() has no return value.
# Check Cluster optional config values before.
self.assertEqual(cluster.location_id, self.LOCATION_ID)
self.assertIsNone(cluster.state)
self.assertEqual(cluster.serve_nodes, self.SERVE_NODES)
self.assertEqual(cluster.default_storage_type, STORAGE_TYPE_SSD)
# Perform the method and check the result.
result = cluster.reload()
self.assertEqual(result, expected_result)
self.assertEqual(cluster.location_id, LOCATION_ID_FROM_SERVER)
self.assertEqual(cluster.state, STATE)
self.assertEqual(cluster.serve_nodes, SERVE_NODES_FROM_SERVER)
self.assertEqual(cluster.default_storage_type, STORAGE_TYPE_FROM_SERVER)
def test_cluster_create_w_cmek():
import datetime
from google.longrunning import operations_pb2
from google.protobuf.any_pb2 import Any
from google.cloud.bigtable_admin_v2.types import (
bigtable_instance_admin as messages_v2_pb2, )
from google.cloud._helpers import _datetime_to_pb_timestamp
from google.cloud.bigtable.instance import Instance
from google.cloud.bigtable_admin_v2.types import instance as instance_v2_pb2
from google.cloud.bigtable.enums import StorageType
NOW = datetime.datetime.utcnow()
NOW_PB = _datetime_to_pb_timestamp(NOW)
credentials = _make_credentials()
client = _make_client(project=PROJECT, credentials=credentials, admin=True)
STORAGE_TYPE_SSD = StorageType.SSD
LOCATION = LOCATION_PATH + LOCATION_ID
instance = Instance(INSTANCE_ID, client)
cluster = _make_cluster(
CLUSTER_ID,
instance,
location_id=LOCATION_ID,
serve_nodes=SERVE_NODES,
default_storage_type=STORAGE_TYPE_SSD,
kms_key_name=KMS_KEY_NAME,
)
name = instance.name
metadata = messages_v2_pb2.CreateClusterMetadata(request_time=NOW_PB)
type_url = "type.googleapis.com/{}".format(
messages_v2_pb2.CreateClusterMetadata._meta._pb.DESCRIPTOR.full_name)
response_pb = operations_pb2.Operation(
name=OP_NAME,
metadata=Any(type_url=type_url,
value=metadata._pb.SerializeToString()),
)
api = client._instance_admin_client = _make_instance_admin_client()
api.common_location_path.return_value = LOCATION
api.instance_path.return_value = name
api.create_cluster.return_value = response_pb
cluster.create()
expected_request_cluster = instance_v2_pb2.Cluster(
location=LOCATION,
serve_nodes=cluster.serve_nodes,
default_storage_type=cluster.default_storage_type,
encryption_config=instance_v2_pb2.Cluster.EncryptionConfig(
kms_key_name=KMS_KEY_NAME, ),
)
expected_request = {
"parent": instance.name,
"cluster_id": CLUSTER_ID,
"cluster": expected_request_cluster,
}
api.create_cluster.assert_called_once_with(request=expected_request)
def test_cluster_from_pb_w_bad_cluster_name():
from google.cloud.bigtable_admin_v2.types import instance as data_v2_pb2
from google.cloud.bigtable.cluster import Cluster
bad_cluster_name = "BAD_NAME"
cluster_pb = data_v2_pb2.Cluster(name=bad_cluster_name)
with pytest.raises(ValueError):
Cluster.from_pb(cluster_pb, None)
def test_from_pb_bad_cluster_name(self):
from google.cloud.bigtable_admin_v2.types import instance as data_v2_pb2
bad_cluster_name = "BAD_NAME"
cluster_pb = data_v2_pb2.Cluster(name=bad_cluster_name)
klass = self._get_target_class()
with self.assertRaises(ValueError):
klass.from_pb(cluster_pb, None)
def test_cluster_reload():
from google.cloud.bigtable_admin_v2.types import instance as data_v2_pb2
from google.cloud.bigtable.enums import StorageType
from google.cloud.bigtable.enums import Cluster
credentials = _make_credentials()
client = _make_client(project=PROJECT, credentials=credentials, admin=True)
STORAGE_TYPE_SSD = StorageType.SSD
instance = _Instance(INSTANCE_ID, client)
cluster = _make_cluster(
CLUSTER_ID,
instance,
location_id=LOCATION_ID,
serve_nodes=SERVE_NODES,
default_storage_type=STORAGE_TYPE_SSD,
kms_key_name=KMS_KEY_NAME,
)
# Create response_pb
LOCATION_ID_FROM_SERVER = "new-location-id"
STATE = Cluster.State.READY
SERVE_NODES_FROM_SERVER = 10
STORAGE_TYPE_FROM_SERVER = StorageType.HDD
response_pb = data_v2_pb2.Cluster(
name=cluster.name,
location=LOCATION_PATH + LOCATION_ID_FROM_SERVER,
state=STATE,
serve_nodes=SERVE_NODES_FROM_SERVER,
default_storage_type=STORAGE_TYPE_FROM_SERVER,
)
# Patch the stub used by the API method.
api = client._instance_admin_client = _make_instance_admin_client()
api.get_cluster.side_effect = [response_pb]
# Create expected_result.
expected_result = None # reload() has no return value.
# Check Cluster optional config values before.
assert cluster.location_id == LOCATION_ID
assert cluster.state is None
assert cluster.serve_nodes == SERVE_NODES
assert cluster.default_storage_type == STORAGE_TYPE_SSD
# Perform the method and check the result.
result = cluster.reload()
assert result == expected_result
assert cluster.location_id == LOCATION_ID_FROM_SERVER
assert cluster.state == STATE
assert cluster.serve_nodes == SERVE_NODES_FROM_SERVER
assert cluster.default_storage_type == STORAGE_TYPE_FROM_SERVER
assert cluster.kms_key_name is None
api.get_cluster.assert_called_once_with(request={"name": cluster.name})
def test_instance_list_clusters():
from google.cloud.bigtable_admin_v2.types import (
bigtable_instance_admin as messages_v2_pb2, )
from google.cloud.bigtable_admin_v2.types import instance as data_v2_pb2
from google.cloud.bigtable.instance import Instance
from google.cloud.bigtable.instance import Cluster
credentials = _make_credentials()
client = _make_client(project=PROJECT, credentials=credentials, admin=True)
instance = Instance(INSTANCE_ID, client)
failed_location = "FAILED"
cluster_id1 = "cluster-id1"
cluster_id2 = "cluster-id2"
cluster_path_template = "projects/{}/instances/{}/clusters/{}"
cluster_name1 = cluster_path_template.format(PROJECT, INSTANCE_ID,
cluster_id1)
cluster_name2 = cluster_path_template.format(PROJECT, INSTANCE_ID,
cluster_id2)
response_pb = messages_v2_pb2.ListClustersResponse(
failed_locations=[failed_location],
clusters=[
data_v2_pb2.Cluster(name=cluster_name1),
data_v2_pb2.Cluster(name=cluster_name2),
],
)
api = client._instance_admin_client = _make_instance_admin_api()
api.list_clusters.side_effect = [response_pb]
api.cluster_path = cluster_path_template.format
# Perform the method and check the result.
clusters, failed_locations = instance.list_clusters()
cluster_1, cluster_2 = clusters
assert isinstance(cluster_1, Cluster)
assert cluster_1.name == cluster_name1
assert isinstance(cluster_2, Cluster)
assert cluster_2.name == cluster_name2
assert failed_locations == [failed_location]
def _to_pb(self):
""" Create cluster proto buff message for API calls """
client = self._instance._client
location = client.instance_admin_client.common_location_path(
client.project, self.location_id)
cluster_pb = instance.Cluster(
location=location,
serve_nodes=self.serve_nodes,
default_storage_type=self.default_storage_type,
)
return cluster_pb
def test_from_pb_project_mistmatch(self):
from google.cloud.bigtable_admin_v2.types import instance as data_v2_pb2
ALT_PROJECT = "ALT_PROJECT"
client = _Client(project=ALT_PROJECT)
instance = _Instance(self.INSTANCE_ID, client)
self.assertNotEqual(self.PROJECT, ALT_PROJECT)
cluster_pb = data_v2_pb2.Cluster(name=self.CLUSTER_NAME)
klass = self._get_target_class()
with self.assertRaises(ValueError):
klass.from_pb(cluster_pb, instance)
def test_cluster_from_pb_w_project_mistmatch():
from google.cloud.bigtable_admin_v2.types import instance as data_v2_pb2
from google.cloud.bigtable.cluster import Cluster
ALT_PROJECT = "ALT_PROJECT"
client = _Client(project=ALT_PROJECT)
instance = _Instance(INSTANCE_ID, client)
assert PROJECT != ALT_PROJECT
cluster_pb = data_v2_pb2.Cluster(name=CLUSTER_NAME)
with pytest.raises(ValueError):
Cluster.from_pb(cluster_pb, instance)
def test_cluster_exists_hit():
from google.cloud.bigtable_admin_v2.types import instance as data_v2_pb2
from google.cloud.bigtable.instance import Instance
credentials = _make_credentials()
client = _make_client(project=PROJECT, credentials=credentials, admin=True)
instance = Instance(INSTANCE_ID, client)
cluster_name = client.instance_admin_client.cluster_path(
PROJECT, INSTANCE_ID, CLUSTER_ID)
response_pb = data_v2_pb2.Cluster(name=cluster_name)
api = client._instance_admin_client = _make_instance_admin_client()
api.get_cluster.return_value = response_pb
cluster = _make_cluster(CLUSTER_ID, instance)
assert cluster.exists()
api.get_cluster.assert_called_once_with(request={"name": cluster.name})
def test_exists(self):
from google.cloud.bigtable_admin_v2.services.bigtable_instance_admin import (
BigtableInstanceAdminClient,
)
from google.cloud.bigtable_admin_v2.types import instance as data_v2_pb2
from google.cloud.bigtable.instance import Instance
from google.api_core import exceptions
instance_api = mock.create_autospec(BigtableInstanceAdminClient)
credentials = _make_credentials()
client = self._make_client(
project=self.PROJECT, credentials=credentials, admin=True
)
instance = Instance(self.INSTANCE_ID, client)
# Create response_pb
cluster_name = client.instance_admin_client.cluster_path(
self.PROJECT, self.INSTANCE_ID, self.CLUSTER_ID
)
response_pb = data_v2_pb2.Cluster(name=cluster_name)
# Patch the stub used by the API method.
client._instance_admin_client = instance_api
bigtable_instance_stub = client._instance_admin_client
bigtable_instance_stub.get_cluster.side_effect = [
response_pb,
exceptions.NotFound("testing"),
exceptions.BadRequest("testing"),
]
# Perform the method and check the result.
non_existing_cluster_id = "cluster-id-2"
alt_cluster_1 = self._make_one(self.CLUSTER_ID, instance)
alt_cluster_2 = self._make_one(non_existing_cluster_id, instance)
self.assertTrue(alt_cluster_1.exists())
self.assertFalse(alt_cluster_2.exists())
with self.assertRaises(exceptions.BadRequest):
alt_cluster_1.exists()
def test_bigtable_cluster_from_pb():
# [START bigtable_api_cluster_from_pb]
from google.cloud.bigtable import Client
from google.cloud.bigtable_admin_v2.types import instance as data_v2_pb2
client = Client(admin=True)
instance = client.instance(INSTANCE_ID)
cluster = instance.cluster(CLUSTER_ID)
name = cluster.name
cluster_state = cluster.state
serve_nodes = 1
cluster_pb = data_v2_pb2.Cluster(
name=name,
location=LOCATION_ID,
state=cluster_state,
serve_nodes=serve_nodes,
default_storage_type=STORAGE_TYPE,
)
cluster2 = cluster.from_pb(cluster_pb, instance)
# [END bigtable_api_cluster_from_pb]
assert cluster2.name == cluster.name
def test_create(self):
import datetime
from google.longrunning import operations_pb2
from google.protobuf.any_pb2 import Any
from google.cloud.bigtable_admin_v2.types import (
bigtable_instance_admin as messages_v2_pb2,
)
from google.cloud._helpers import _datetime_to_pb_timestamp
from google.cloud.bigtable.instance import Instance
from google.cloud.bigtable_admin_v2.services.bigtable_instance_admin import (
BigtableInstanceAdminClient,
)
from google.cloud.bigtable_admin_v2.types import instance as instance_v2_pb2
from google.cloud.bigtable.enums import StorageType
NOW = datetime.datetime.utcnow()
NOW_PB = _datetime_to_pb_timestamp(NOW)
credentials = _make_credentials()
client = self._make_client(
project=self.PROJECT, credentials=credentials, admin=True
)
STORAGE_TYPE_SSD = StorageType.SSD
LOCATION = self.LOCATION_PATH + self.LOCATION_ID
instance = Instance(self.INSTANCE_ID, client)
cluster = self._make_one(
self.CLUSTER_ID,
instance,
location_id=self.LOCATION_ID,
serve_nodes=self.SERVE_NODES,
default_storage_type=STORAGE_TYPE_SSD,
)
expected_request_cluster = instance_v2_pb2.Cluster(
location=LOCATION,
serve_nodes=cluster.serve_nodes,
default_storage_type=cluster.default_storage_type,
)
expected_request = {
"request": {
"parent": instance.name,
"cluster_id": self.CLUSTER_ID,
"cluster": expected_request_cluster,
}
}
name = instance.name
metadata = messages_v2_pb2.CreateClusterMetadata(request_time=NOW_PB)
type_url = "type.googleapis.com/{}".format(
messages_v2_pb2.CreateClusterMetadata._meta._pb.DESCRIPTOR.full_name
)
response_pb = operations_pb2.Operation(
name=self.OP_NAME,
metadata=Any(type_url=type_url, value=metadata._pb.SerializeToString()),
)
# Patch the stub used by the API method.
api = mock.create_autospec(BigtableInstanceAdminClient)
api.common_location_path.return_value = LOCATION
client._instance_admin_client = api
cluster._instance._client = client
cluster._instance._client.instance_admin_client.instance_path.return_value = (
name
)
client._instance_admin_client.create_cluster.return_value = response_pb
# Perform the method and check the result.
cluster.create()
actual_request = client._instance_admin_client.create_cluster.call_args_list[
0
].kwargs
self.assertEqual(actual_request, expected_request)
def test_list_clusters(self):
from google.cloud.bigtable_admin_v2.services.bigtable_instance_admin import (
BigtableInstanceAdminClient, )
from google.cloud.bigtable_admin_v2.types import (
bigtable_instance_admin as messages_v2_pb2, )
from google.cloud.bigtable_admin_v2.types import instance as data_v2_pb2
from google.cloud.bigtable.instance import Cluster
instance_api = mock.create_autospec(BigtableInstanceAdminClient)
credentials = _make_credentials()
client = self._make_one(project=self.PROJECT,
credentials=credentials,
admin=True)
INSTANCE_ID1 = "instance-id1"
INSTANCE_ID2 = "instance-id2"
failed_location = "FAILED"
cluster_id1 = "{}-cluster".format(INSTANCE_ID1)
cluster_id2 = "{}-cluster-1".format(INSTANCE_ID2)
cluster_id3 = "{}-cluster-2".format(INSTANCE_ID2)
cluster_name1 = client.instance_admin_client.cluster_path(
self.PROJECT, INSTANCE_ID1, cluster_id1)
cluster_name2 = client.instance_admin_client.cluster_path(
self.PROJECT, INSTANCE_ID2, cluster_id2)
cluster_name3 = client.instance_admin_client.cluster_path(
self.PROJECT, INSTANCE_ID2, cluster_id3)
# Create response_pb
response_pb = messages_v2_pb2.ListClustersResponse(
failed_locations=[failed_location],
clusters=[
data_v2_pb2.Cluster(name=cluster_name1),
data_v2_pb2.Cluster(name=cluster_name2),
data_v2_pb2.Cluster(name=cluster_name3),
],
)
# Patch the stub used by the API method.
client._instance_admin_client = instance_api
instance_stub = client._instance_admin_client
instance_stub.list_clusters.side_effect = [response_pb]
# Perform the method and check the result.
clusters, failed_locations = client.list_clusters()
cluster_1, cluster_2, cluster_3 = clusters
self.assertIsInstance(cluster_1, Cluster)
self.assertEqual(cluster_1.cluster_id, cluster_id1)
self.assertEqual(cluster_1._instance.instance_id, INSTANCE_ID1)
self.assertIsInstance(cluster_2, Cluster)
self.assertEqual(cluster_2.cluster_id, cluster_id2)
self.assertEqual(cluster_2._instance.instance_id, INSTANCE_ID2)
self.assertIsInstance(cluster_3, Cluster)
self.assertEqual(cluster_3.cluster_id, cluster_id3)
self.assertEqual(cluster_3._instance.instance_id, INSTANCE_ID2)
self.assertEqual(failed_locations, [failed_location])
