def test_bigtable_list_instances():
# [START bigtable_list_instances]
from google.cloud.bigtable import Client
client = Client(admin=True)
(instances_list, failed_locations_list) = client.list_instances()
# [END bigtable_list_instances]
assert len(instances_list) > 0
class BigtableIOWriteTest(unittest.TestCase):
""" Bigtable Write Connector Test
"""
DEFAULT_TABLE_PREFIX = "python-test"
instance_id = DEFAULT_TABLE_PREFIX + "-" + str(uuid.uuid4())[:8]
cluster_id = DEFAULT_TABLE_PREFIX + "-" + str(uuid.uuid4())[:8]
table_id = DEFAULT_TABLE_PREFIX + "-" + str(uuid.uuid4())[:8]
number = 500
LOCATION_ID = "us-east1-b"
def setUp(self):
try:
from google.cloud.bigtable import enums
self.STORAGE_TYPE = enums.StorageType.HDD
self.INSTANCE_TYPE = enums.Instance.Type.DEVELOPMENT
except ImportError:
self.STORAGE_TYPE = 2
self.INSTANCE_TYPE = 2
self.test_pipeline = TestPipeline(is_integration_test=True)
self.runner_name = type(self.test_pipeline.runner).__name__
self.project = self.test_pipeline.get_option('project')
self.client = Client(project=self.project, admin=True)
self._delete_old_instances()
self.instance = self.client.instance(self.instance_id,
instance_type=self.INSTANCE_TYPE,
labels=LABELS)
if not self.instance.exists():
cluster = self.instance.cluster(
self.cluster_id,
self.LOCATION_ID,
default_storage_type=self.STORAGE_TYPE)
self.instance.create(clusters=[cluster])
self.table = self.instance.table(self.table_id)
if not self.table.exists():
max_versions_rule = column_family.MaxVersionsGCRule(2)
column_family_id = 'cf1'
column_families = {column_family_id: max_versions_rule}
self.table.create(column_families=column_families)
def _delete_old_instances(self):
instances = self.client.list_instances()
EXISTING_INSTANCES[:] = instances
def age_in_hours(micros):
return (
datetime.datetime.utcnow().replace(tzinfo=UTC) -
(_datetime_from_microseconds(micros))).total_seconds() // 3600
CLEAN_INSTANCE = [
i for instance in EXISTING_INSTANCES for i in instance
if (LABEL_KEY in i.labels.keys() and (
age_in_hours(int(i.labels[LABEL_KEY])) >= 2))
]
if CLEAN_INSTANCE:
for instance in CLEAN_INSTANCE:
instance.delete()
def tearDown(self):
if self.instance.exists():
self.instance.delete()
def test_bigtable_write(self):
number = self.number
pipeline_args = self.test_pipeline.options_list
pipeline_options = PipelineOptions(pipeline_args)
with beam.Pipeline(options=pipeline_options) as pipeline:
config_data = {
'project_id': self.project,
'instance_id': self.instance,
'table_id': self.table
}
_ = (pipeline
| 'Generate Direct Rows' >> GenerateTestRows(
number, **config_data))
assert pipeline.result.state == PipelineState.DONE
read_rows = self.table.read_rows()
assert len([_ for _ in read_rows]) == number
if not hasattr(pipeline.result, 'has_job') or pipeline.result.has_job:
read_filter = MetricsFilter().with_name('Written Row')
query_result = pipeline.result.metrics().query(read_filter)
if query_result['counters']:
read_counter = query_result['counters'][0]
logging.info('Number of Rows: %d', read_counter.committed)
assert read_counter.committed == number
