def test_do_mutate_retryable_rows(self):
        from google.cloud.bigtable.row import DirectRow
        from tests.unit._testing import _FakeStub

        # Setup:
        #   - Mutate 2 rows.
        # Action:
        #   - Initial attempt will mutate all 2 rows.
        # Expectation:
        #   - Expect [success, non-retryable]

        client = _Client()
        instance = _Instance(self.INSTANCE_NAME, client=client)
        table = self._make_table(self.TABLE_ID, instance)

        row_1 = DirectRow(row_key=b'row_key', table=table)
        row_1.set_cell('cf', b'col', b'value1')
        row_2 = DirectRow(row_key=b'row_key_2', table=table)
        row_2.set_cell('cf', b'col', b'value2')

        response = self._make_responses([self.SUCCESS, self.NON_RETRYABLE])

        # Patch the stub used by the API method.
        client._data_stub = _FakeStub([response])

        worker = self._make_worker(table._instance._client, table.name,
                                   [row_1, row_2])
        statuses = worker._do_mutate_retryable_rows()

        result = [status.code for status in statuses]
        expected_result = [self.SUCCESS, self.NON_RETRYABLE]

        self.assertEqual(result, expected_result)
Esempio n. 2
0
    def test_do_mutate_retryable_rows_mismatch_num_responses(self):
        from google.cloud.bigtable._generated.bigtable_pb2 import (
            MutateRowsResponse)
        from google.cloud.bigtable.row import DirectRow
        from google.rpc.status_pb2 import Status
        from tests.unit._testing import _FakeStub

        client = _Client()
        instance = _Instance(self.INSTANCE_NAME, client=client)
        table = self._make_table(self.TABLE_ID, instance)

        row_1 = DirectRow(row_key=b'row_key', table=table)
        row_1.set_cell('cf', b'col', b'value1')
        row_2 = DirectRow(row_key=b'row_key_2', table=table)
        row_2.set_cell('cf', b'col', b'value2')

        response = MutateRowsResponse(entries=[
            MutateRowsResponse.Entry(
                index=0,
                status=Status(code=0),
            ),
        ], )

        # Patch the stub used by the API method.
        client._data_stub = _FakeStub([response])

        worker = self._make_worker(table._instance._client, table.name,
                                   [row_1, row_2])
        with self.assertRaises(RuntimeError):
            worker._do_mutate_retryable_rows()
Esempio n. 3
0
    def test_delete(self):
        from google.protobuf import empty_pb2
        from tests.unit._testing import _FakeStub

        client = _Client()
        instance = _Instance(self.INSTANCE_NAME, client=client)
        table = self._make_one(self.TABLE_ID, instance)

        # Create request_pb
        request_pb = _DeleteTableRequestPB(name=self.TABLE_NAME)

        # Create response_pb
        response_pb = empty_pb2.Empty()

        # Patch the stub used by the API method.
        client._table_stub = stub = _FakeStub(response_pb)

        # Create expected_result.
        expected_result = None  # delete() has no return value.

        # Perform the method and check the result.
        result = table.delete()
        self.assertEqual(result, expected_result)
        self.assertEqual(stub.method_calls, [(
            'DeleteTable',
            (request_pb,),
            {},
        )])
Esempio n. 4
0
    def test_do_mutate_retryable_rows(self):
        from google.cloud.bigtable.row import DirectRow
        from tests.unit._testing import _FakeStub

        # Setup:
        #   - Mutate 2 rows.
        # Action:
        #   - Initial attempt will mutate all 2 rows.
        # Expectation:
        #   - Expect [success, non-retryable]

        client = _Client()
        instance = _Instance(self.INSTANCE_NAME, client=client)
        table = self._make_table(self.TABLE_ID, instance)

        row_1 = DirectRow(row_key=b'row_key', table=table)
        row_1.set_cell('cf', b'col', b'value1')
        row_2 = DirectRow(row_key=b'row_key_2', table=table)
        row_2.set_cell('cf', b'col', b'value2')

        response = self._make_responses([self.SUCCESS, self.NON_RETRYABLE])

        # Patch the stub used by the API method.
        client._data_stub = _FakeStub([response])

        worker = self._make_worker(
            table._instance._client, table.name, [row_1, row_2])
        statuses = worker._do_mutate_retryable_rows()

        result = [status.code for status in statuses]
        expected_result = [self.SUCCESS, self.NON_RETRYABLE]

        self.assertEqual(result, expected_result)
Esempio n. 5
0
    def test_sample_row_keys(self):
        from tests.unit._testing import _FakeStub

        client = _Client()
        instance = _Instance(self.INSTANCE_NAME, client=client)
        table = self._make_one(self.TABLE_ID, instance)

        # Create request_pb
        request_pb = _SampleRowKeysRequestPB(table_name=self.TABLE_NAME)

        # Create response_iterator
        response_iterator = object()  # Just passed to a mock.

        # Patch the stub used by the API method.
        client._data_stub = stub = _FakeStub(response_iterator)

        # Create expected_result.
        expected_result = response_iterator

        # Perform the method and check the result.
        result = table.sample_row_keys()
        self.assertEqual(result, expected_result)
        self.assertEqual(stub.method_calls, [(
            'SampleRowKeys',
            (request_pb,),
            {},
        )])
Esempio n. 6
0
    def test_yield_rows(self):
        from tests.unit._testing import _FakeStub

        client = _Client()
        instance = _Instance(self.INSTANCE_NAME, client=client)
        table = self._make_one(self.TABLE_ID, instance)

        # Create response_iterator
        chunk = _ReadRowsResponseCellChunkPB(
            row_key=self.ROW_KEY,
            family_name=self.FAMILY_NAME,
            qualifier=self.QUALIFIER,
            timestamp_micros=self.TIMESTAMP_MICROS,
            value=self.VALUE,
            commit_row=True,
        )
        chunks = [chunk]

        response = _ReadRowsResponseV2(chunks)
        response_iterator = _MockCancellableIterator(response)

        # Patch the stub used by the API method.
        client._data_stub = _FakeStub(response_iterator)

        rows = []
        for row in table.yield_rows():
            rows.append(row)
        result = rows[0]

        self.assertEqual(result.row_key, self.ROW_KEY)
Esempio n. 7
0
    def test_create_w_explicit_serve_nodes(self):
        from google.api_core import operation
        from google.longrunning import operations_pb2
        from tests.unit._testing import _FakeStub
        from google.cloud.bigtable_admin_v2.gapic import (
            bigtable_instance_admin_client)

        api = bigtable_instance_admin_client.BigtableInstanceAdminClient(
            mock.Mock())
        credentials = _make_credentials()
        client = self._make_client(project=self.PROJECT,
                                   credentials=credentials,
                                   admin=True)
        instance = self._make_one(self.INSTANCE_ID, client, self.LOCATION_ID)

        # Create response_pb
        response_pb = operations_pb2.Operation(name=self.OP_NAME)

        # Patch the stub used by the API method.
        stub = _FakeStub(response_pb)
        client._instance_admin_client = api
        client._instance_admin_client.bigtable_instance_admin_stub = stub

        # Perform the method and check the result.
        result = instance.create()

        self.assertIsInstance(result, operation.Operation)
    def test_update(self):
        from google.cloud.bigtable._generated import (instance_pb2 as
                                                      data_v2_pb2)
        from tests.unit._testing import _FakeStub

        client = _Client(self.PROJECT)
        instance = self._make_one(self.INSTANCE_ID,
                                  client,
                                  self.LOCATION_ID,
                                  display_name=self.DISPLAY_NAME)

        # Create request_pb
        request_pb = data_v2_pb2.Instance(
            name=self.INSTANCE_NAME,
            display_name=self.DISPLAY_NAME,
        )

        # Create response_pb
        response_pb = data_v2_pb2.Instance()

        # Patch the stub used by the API method.
        client._instance_stub = stub = _FakeStub(response_pb)

        # Create expected_result.
        expected_result = None

        # Perform the method and check the result.
        result = instance.update()

        self.assertEqual(result, expected_result)
        self.assertEqual(stub.method_calls, [(
            'UpdateInstance',
            (request_pb, ),
            {},
        )])
    def test_delete(self):
        from google.protobuf import empty_pb2
        from tests.unit._testing import _FakeStub

        client = _Client()
        instance = _Instance(self.INSTANCE_NAME, client=client)
        table = self._make_one(self.TABLE_ID, instance)

        # Create request_pb
        request_pb = _DeleteTableRequestPB(name=self.TABLE_NAME)

        # Create response_pb
        response_pb = empty_pb2.Empty()

        # Patch the stub used by the API method.
        client._table_stub = stub = _FakeStub(response_pb)

        # Create expected_result.
        expected_result = None  # delete() has no return value.

        # Perform the method and check the result.
        result = table.delete()
        self.assertEqual(result, expected_result)
        self.assertEqual(stub.method_calls, [(
            'DeleteTable',
            (request_pb, ),
            {},
        )])
Esempio n. 10
0
    def _list_column_families_helper(self):
        from tests.unit._testing import _FakeStub

        client = _Client()
        instance = _Instance(self.INSTANCE_NAME, client=client)
        table = self._make_one(self.TABLE_ID, instance)

        # Create request_pb
        request_pb = _GetTableRequestPB(name=self.TABLE_NAME)

        # Create response_pb
        COLUMN_FAMILY_ID = 'foo'
        column_family = _ColumnFamilyPB()
        response_pb = _TablePB(
            column_families={COLUMN_FAMILY_ID: column_family}, )

        # Patch the stub used by the API method.
        client._table_stub = stub = _FakeStub(response_pb)

        # Create expected_result.
        expected_result = {
            COLUMN_FAMILY_ID: table.column_family(COLUMN_FAMILY_ID),
        }

        # Perform the method and check the result.
        result = table.list_column_families()
        self.assertEqual(result, expected_result)
        self.assertEqual(stub.method_calls, [(
            'GetTable',
            (request_pb, ),
            {},
        )])
Esempio n. 11
0
    def test_delete(self):
        from google.protobuf import empty_pb2
        from google.cloud.bigtable._generated import (
            bigtable_instance_admin_pb2 as messages_v2_pb)
        from tests.unit._testing import _FakeStub

        client = _Client(self.PROJECT)
        instance = self._make_one(self.INSTANCE_ID, client, self.LOCATION_ID)

        # Create request_pb
        request_pb = messages_v2_pb.DeleteInstanceRequest(
            name=self.INSTANCE_NAME)

        # Create response_pb
        response_pb = empty_pb2.Empty()

        # Patch the stub used by the API method.
        client._instance_stub = stub = _FakeStub(response_pb)

        # Create expected_result.
        expected_result = None  # delete() has no return value.

        # Perform the method and check the result.
        result = instance.delete()

        self.assertEqual(result, expected_result)
        self.assertEqual(stub.method_calls, [(
            'DeleteInstance',
            (request_pb,),
            {},
        )])
    def test_delete(self):
        from google.protobuf import empty_pb2
        from google.cloud.bigtable._generated import (
            bigtable_instance_admin_pb2 as messages_v2_pb)
        from tests.unit._testing import _FakeStub

        client = _Client(self.PROJECT)
        instance = self._make_one(self.INSTANCE_ID, client, self.LOCATION_ID)

        # Create request_pb
        request_pb = messages_v2_pb.DeleteInstanceRequest(
            name=self.INSTANCE_NAME)

        # Create response_pb
        response_pb = empty_pb2.Empty()

        # Patch the stub used by the API method.
        client._instance_stub = stub = _FakeStub(response_pb)

        # Create expected_result.
        expected_result = None  # delete() has no return value.

        # Perform the method and check the result.
        result = instance.delete()

        self.assertEqual(result, expected_result)
        self.assertEqual(stub.method_calls, [(
            'DeleteInstance',
            (request_pb, ),
            {},
        )])
Esempio n. 13
0
    def test_update(self):
        from google.cloud.bigtable._generated import (
            instance_pb2 as data_v2_pb2)
        from tests.unit._testing import _FakeStub

        client = _Client(self.PROJECT)
        instance = self._make_one(self.INSTANCE_ID, client, self.LOCATION_ID,
                                  display_name=self.DISPLAY_NAME)

        # Create request_pb
        request_pb = data_v2_pb2.Instance(
            name=self.INSTANCE_NAME,
            display_name=self.DISPLAY_NAME,
        )

        # Create response_pb
        response_pb = data_v2_pb2.Instance()

        # Patch the stub used by the API method.
        client._instance_stub = stub = _FakeStub(response_pb)

        # Create expected_result.
        expected_result = None

        # Perform the method and check the result.
        result = instance.update()

        self.assertEqual(result, expected_result)
        self.assertEqual(stub.method_calls, [(
            'UpdateInstance',
            (request_pb,),
            {},
        )])
Esempio n. 14
0
    def test_mutate_rows(self):
        from google.cloud.bigtable._generated.bigtable_pb2 import (
            MutateRowsResponse)
        from google.cloud.bigtable.row import DirectRow
        from google.rpc.status_pb2 import Status
        from tests.unit._testing import _FakeStub

        client = _Client()
        instance = _Instance(self.INSTANCE_NAME, client=client)
        table = self._make_one(self.TABLE_ID, instance)

        row_1 = DirectRow(row_key=b'row_key', table=table)
        row_1.set_cell('cf', b'col', b'value1')
        row_2 = DirectRow(row_key=b'row_key_2', table=table)
        row_2.set_cell('cf', b'col', b'value2')

        response = MutateRowsResponse(entries=[
            MutateRowsResponse.Entry(
                index=0,
                status=Status(code=0),
            ),
            MutateRowsResponse.Entry(
                index=1,
                status=Status(code=1),
            ),
        ], )

        # Patch the stub used by the API method.
        client._data_stub = _FakeStub([response])
        statuses = table.mutate_rows([row_1, row_2])
        result = [status.code for status in statuses]
        expected_result = [0, 1]

        self.assertEqual(result, expected_result)
Esempio n. 15
0
    def test_sample_row_keys(self):
        from tests.unit._testing import _FakeStub

        client = _Client()
        instance = _Instance(self.INSTANCE_NAME, client=client)
        table = self._make_one(self.TABLE_ID, instance)

        # Create request_pb
        request_pb = _SampleRowKeysRequestPB(table_name=self.TABLE_NAME)

        # Create response_iterator
        response_iterator = object()  # Just passed to a mock.

        # Patch the stub used by the API method.
        client._data_stub = stub = _FakeStub(response_iterator)

        # Create expected_result.
        expected_result = response_iterator

        # Perform the method and check the result.
        result = table.sample_row_keys()
        self.assertEqual(result, expected_result)
        self.assertEqual(stub.method_calls, [(
            'SampleRowKeys',
            (request_pb, ),
            {},
        )])
Esempio n. 16
0
    def _list_column_families_helper(self):
        from tests.unit._testing import _FakeStub

        client = _Client()
        instance = _Instance(self.INSTANCE_NAME, client=client)
        table = self._make_one(self.TABLE_ID, instance)

        # Create request_pb
        request_pb = _GetTableRequestPB(name=self.TABLE_NAME)

        # Create response_pb
        COLUMN_FAMILY_ID = 'foo'
        column_family = _ColumnFamilyPB()
        response_pb = _TablePB(
            column_families={COLUMN_FAMILY_ID: column_family},
        )

        # Patch the stub used by the API method.
        client._table_stub = stub = _FakeStub(response_pb)

        # Create expected_result.
        expected_result = {
            COLUMN_FAMILY_ID: table.column_family(COLUMN_FAMILY_ID),
        }

        # Perform the method and check the result.
        result = table.list_column_families()
        self.assertEqual(result, expected_result)
        self.assertEqual(stub.method_calls, [(
            'GetTable',
            (request_pb,),
            {},
        )])
Esempio n. 17
0
    def test_do_mutate_retryable_rows_second_retry(self):
        from google.cloud.bigtable.row import DirectRow
        from google.cloud.bigtable.table import _BigtableRetryableError
        from tests.unit._testing import _FakeStub

        # Setup:
        #   - Mutate 4 rows.
        #   - First try results:
        #       [success, retryable, non-retryable, retryable]
        # Action:
        #   - Second try should re-attempt the 'retryable' rows.
        # Expectation:
        #   - After second try:
        #       [success, success, non-retryable, retryable]
        #   - One of the rows tried second time returns retryable error code,
        #     so expect a raise.
        #   - Exception contains response whose index should be '3' even though
        #     only two rows were retried.

        client = _Client()
        instance = _Instance(self.INSTANCE_NAME, client=client)
        table = self._make_table(self.TABLE_ID, instance)

        row_1 = DirectRow(row_key=b'row_key', table=table)
        row_1.set_cell('cf', b'col', b'value1')
        row_2 = DirectRow(row_key=b'row_key_2', table=table)
        row_2.set_cell('cf', b'col', b'value2')
        row_3 = DirectRow(row_key=b'row_key_3', table=table)
        row_3.set_cell('cf', b'col', b'value3')
        row_4 = DirectRow(row_key=b'row_key_4', table=table)
        row_4.set_cell('cf', b'col', b'value4')

        response = self._make_responses([self.SUCCESS, self.RETRYABLE_1])

        # Patch the stub used by the API method.
        client._data_stub = _FakeStub([response])

        worker = self._make_worker(
            table._instance._client,
            table.name, [row_1, row_2, row_3, row_4])
        worker.responses_statuses = self._make_responses_statuses([
            self.SUCCESS,
            self.RETRYABLE_1,
            self.NON_RETRYABLE,
            self.RETRYABLE_2])

        with self.assertRaises(_BigtableRetryableError):
            worker._do_mutate_retryable_rows()

        statuses = worker.responses_statuses
        result = [status.code for status in statuses]
        expected_result = [self.SUCCESS,
                           self.SUCCESS,
                           self.NON_RETRYABLE,
                           self.RETRYABLE_1]

        self.assertEqual(result, expected_result)
Esempio n. 18
0
    def test_create(self):
        import datetime
        from google.longrunning import operations_pb2
        from google.protobuf.any_pb2 import Any
        from google.cloud.bigtable._generated import (
            bigtable_instance_admin_pb2 as messages_v2_pb2)
        from google.cloud._helpers import _datetime_to_pb_timestamp
        from tests.unit._testing import _FakeStub
        from google.cloud.operation import Operation
        from google.cloud.bigtable.cluster import DEFAULT_SERVE_NODES

        NOW = datetime.datetime.utcnow()
        NOW_PB = _datetime_to_pb_timestamp(NOW)
        client = _Client(self.PROJECT)
        instance = self._make_one(self.INSTANCE_ID, client, self.LOCATION_ID,
                                  display_name=self.DISPLAY_NAME)

        # Create response_pb
        metadata = messages_v2_pb2.CreateInstanceMetadata(request_time=NOW_PB)
        type_url = 'type.googleapis.com/%s' % (
            messages_v2_pb2.CreateInstanceMetadata.DESCRIPTOR.full_name,)
        response_pb = operations_pb2.Operation(
            name=self.OP_NAME,
            metadata=Any(
                type_url=type_url,
                value=metadata.SerializeToString(),
            )
        )

        # Patch the stub used by the API method.
        client._instance_stub = stub = _FakeStub(response_pb)

        # Perform the method and check the result.
        result = instance.create()

        self.assertIsInstance(result, Operation)
        self.assertEqual(result.name, self.OP_NAME)
        self.assertIs(result.target, instance)
        self.assertIs(result.client, client)
        self.assertIsInstance(result.metadata,
                              messages_v2_pb2.CreateInstanceMetadata)
        self.assertEqual(result.metadata.request_time, NOW_PB)
        self.assertEqual(result.caller_metadata,
                         {'request_type': 'CreateInstance'})

        self.assertEqual(len(stub.method_calls), 1)
        api_name, args, kwargs = stub.method_calls[0]
        self.assertEqual(api_name, 'CreateInstance')
        request_pb, = args
        self.assertIsInstance(request_pb,
                              messages_v2_pb2.CreateInstanceRequest)
        self.assertEqual(request_pb.parent, 'projects/%s' % (self.PROJECT,))
        self.assertEqual(request_pb.instance_id, self.INSTANCE_ID)
        self.assertEqual(request_pb.instance.display_name, self.DISPLAY_NAME)
        cluster = request_pb.clusters[self.INSTANCE_ID]
        self.assertEqual(cluster.serve_nodes, DEFAULT_SERVE_NODES)
        self.assertEqual(kwargs, {})
Esempio n. 19
0
    def test_do_mutate_retryable_rows_retry(self):
        from google.api_core.exceptions import ServiceUnavailable
        from google.cloud.bigtable._generated.bigtable_pb2 import (
            MutateRowsResponse)
        from google.cloud.bigtable.row import DirectRow
        from google.rpc.status_pb2 import Status
        from tests.unit._testing import _FakeStub

        # Setup:
        #   - Mutate 3 rows.
        # Action:
        #   - Initial attempt will mutate all 3 rows.
        # Expectation:
        #   - Second row returns retryable error code, so expect a raise.
        #   - State of responses_statuses should be
        #       [success, retryable, non-retryable]

        client = _Client()
        instance = _Instance(self.INSTANCE_NAME, client=client)
        table = self._make_table(self.TABLE_ID, instance)

        row_1 = DirectRow(row_key=b'row_key', table=table)
        row_1.set_cell('cf', b'col', b'value1')
        row_2 = DirectRow(row_key=b'row_key_2', table=table)
        row_2.set_cell('cf', b'col', b'value2')
        row_3 = DirectRow(row_key=b'row_key_3', table=table)
        row_3.set_cell('cf', b'col', b'value3')

        response = MutateRowsResponse(entries=[
            MutateRowsResponse.Entry(
                index=0,
                status=Status(code=0),
            ),
            MutateRowsResponse.Entry(
                index=1,
                status=Status(code=4),
            ),
            MutateRowsResponse.Entry(
                index=2,
                status=Status(code=1),
            ),
        ], )

        # Patch the stub used by the API method.
        client._data_stub = _FakeStub([response])

        worker = self._make_worker(table._instance._client, table.name,
                                   [row_1, row_2, row_3])

        with self.assertRaises(ServiceUnavailable):
            worker._do_mutate_retryable_rows()

        statuses = worker.responses_statuses
        result = [status.code for status in statuses]
        expected_result = [0, 4, 1]

        self.assertEqual(result, expected_result)
Esempio n. 20
0
    def test_list_clusters(self):
        from google.cloud.bigtable._generated import (
            instance_pb2 as instance_v2_pb2)
        from google.cloud.bigtable._generated import (
            bigtable_instance_admin_pb2 as messages_v2_pb2)
        from tests.unit._testing import _FakeStub

        FAILED_LOCATION = 'FAILED'
        FAILED_LOCATIONS = [FAILED_LOCATION]
        CLUSTER_ID1 = 'cluster-id1'
        CLUSTER_ID2 = 'cluster-id2'
        SERVE_NODES = 4

        client = _Client(self.PROJECT)
        instance = self._make_one(self.INSTANCE_ID, client, self.LOCATION_ID)

        CLUSTER_NAME1 = (instance.name + '/clusters/' + CLUSTER_ID1)
        CLUSTER_NAME2 = (instance.name + '/clusters/' + CLUSTER_ID2)
        # Create request_pb
        request_pb = messages_v2_pb2.ListClustersRequest(
            parent=instance.name,
        )

        # Create response_pb
        response_pb = messages_v2_pb2.ListClustersResponse(
            failed_locations=[FAILED_LOCATION],
            clusters=[
                instance_v2_pb2.Cluster(
                    name=CLUSTER_NAME1,
                    serve_nodes=SERVE_NODES,
                ),
                instance_v2_pb2.Cluster(
                    name=CLUSTER_NAME2,
                    serve_nodes=SERVE_NODES,
                ),
            ],
        )

        # Patch the stub used by the API method.
        client._instance_stub = stub = _FakeStub(response_pb)

        # Create expected_result.
        clusters = [
            instance.cluster(CLUSTER_ID1),
            instance.cluster(CLUSTER_ID2),
        ]
        expected_result = (clusters, FAILED_LOCATIONS)

        # Perform the method and check the result.
        result = instance.list_clusters()
        self.assertEqual(result, expected_result)
        self.assertEqual(stub.method_calls, [(
            'ListClusters',
            (request_pb,),
            {},
        )])
Esempio n. 21
0
    def test_list_clusters(self):
        from google.cloud.bigtable._generated import (
            instance_pb2 as instance_v2_pb2)
        from google.cloud.bigtable._generated import (
            bigtable_instance_admin_pb2 as messages_v2_pb2)
        from tests.unit._testing import _FakeStub

        FAILED_LOCATION = 'FAILED'
        FAILED_LOCATIONS = [FAILED_LOCATION]
        CLUSTER_ID1 = 'cluster-id1'
        CLUSTER_ID2 = 'cluster-id2'
        SERVE_NODES = 4

        client = _Client(self.PROJECT)
        instance = self._make_one(self.INSTANCE_ID, client, self.LOCATION_ID)

        CLUSTER_NAME1 = (instance.name + '/clusters/' + CLUSTER_ID1)
        CLUSTER_NAME2 = (instance.name + '/clusters/' + CLUSTER_ID2)
        # Create request_pb
        request_pb = messages_v2_pb2.ListClustersRequest(
            parent=instance.name,
        )

        # Create response_pb
        response_pb = messages_v2_pb2.ListClustersResponse(
            failed_locations=[FAILED_LOCATION],
            clusters=[
                instance_v2_pb2.Cluster(
                    name=CLUSTER_NAME1,
                    serve_nodes=SERVE_NODES,
                ),
                instance_v2_pb2.Cluster(
                    name=CLUSTER_NAME2,
                    serve_nodes=SERVE_NODES,
                ),
            ],
        )

        # Patch the stub used by the API method.
        client._instance_stub = stub = _FakeStub(response_pb)

        # Create expected_result.
        clusters = [
            instance.cluster(CLUSTER_ID1),
            instance.cluster(CLUSTER_ID2),
        ]
        expected_result = (clusters, FAILED_LOCATIONS)

        # Perform the method and check the result.
        result = instance.list_clusters()
        self.assertEqual(result, expected_result)
        self.assertEqual(stub.method_calls, [(
            'ListClusters',
            (request_pb,),
            {},
        )])
Esempio n. 22
0
    def test_commit(self):
        from google.cloud._testing import _Monkey
        from tests.unit._testing import _FakeStub
        from google.cloud.bigtable import row as MUT

        row_key = b'row_key'
        table_name = 'projects/more-stuff'
        column_family_id = u'column_family_id'
        column = b'column'
        client = _Client()
        table = _Table(table_name, client=client)
        row = self._make_one(row_key, table)

        # Create request_pb
        value = b'bytes-value'
        # We will call row.append_cell_value(COLUMN_FAMILY_ID, COLUMN, value).
        request_pb = _ReadModifyWriteRowRequestPB(
            table_name=table_name,
            row_key=row_key,
            rules=[
                _ReadModifyWriteRulePB(
                    family_name=column_family_id,
                    column_qualifier=column,
                    append_value=value,
                ),
            ],
        )

        # Create response_pb
        response_pb = object()

        # Patch the stub used by the API method.
        client._data_stub = stub = _FakeStub(response_pb)

        # Create expected_result.
        row_responses = []
        expected_result = object()

        def mock_parse_rmw_row_response(row_response):
            row_responses.append(row_response)
            return expected_result

        # Perform the method and check the result.
        with _Monkey(MUT, _parse_rmw_row_response=mock_parse_rmw_row_response):
            row.append_cell_value(column_family_id, column, value)
            result = row.commit()

        self.assertEqual(result, expected_result)
        self.assertEqual(stub.method_calls, [(
            'ReadModifyWriteRow',
            (request_pb, ),
            {},
        )])
        self.assertEqual(row_responses, [response_pb])
        self.assertEqual(row._rule_pb_list, [])
Esempio n. 23
0
    def test_do_mutate_retryable_rows_second_try(self):
        from google.cloud.bigtable._generated.bigtable_pb2 import (
            MutateRowsResponse)
        from google.cloud.bigtable.row import DirectRow
        from google.rpc.status_pb2 import Status
        from tests.unit._testing import _FakeStub

        # Setup:
        #   - Mutate 4 rows.
        #   - First try results:
        #       [success, retryable, non-retryable, retryable]
        # Action:
        #   - Second try should re-attempt the 'retryable' rows.
        # Expectation:
        #   - After second try:
        #       [success, non-retryable, non-retryable, success]

        client = _Client()
        instance = _Instance(self.INSTANCE_NAME, client=client)
        table = self._make_table(self.TABLE_ID, instance)

        row_1 = DirectRow(row_key=b'row_key', table=table)
        row_1.set_cell('cf', b'col', b'value1')
        row_2 = DirectRow(row_key=b'row_key_2', table=table)
        row_2.set_cell('cf', b'col', b'value2')
        row_3 = DirectRow(row_key=b'row_key_3', table=table)
        row_3.set_cell('cf', b'col', b'value3')
        row_4 = DirectRow(row_key=b'row_key_4', table=table)
        row_4.set_cell('cf', b'col', b'value4')

        response = MutateRowsResponse(entries=[
            MutateRowsResponse.Entry(
                index=0,
                status=Status(code=1),
            ),
            MutateRowsResponse.Entry(
                index=1,
                status=Status(code=0),
            ),
        ], )

        # Patch the stub used by the API method.
        client._data_stub = _FakeStub([response])

        worker = self._make_worker(table._instance._client, table.name,
                                   [row_1, row_2, row_3, row_4])
        worker.responses_statuses = self._make_responses_statuses(
            [0, 4, 1, 10])

        statuses = worker._do_mutate_retryable_rows()

        result = [status.code for status in statuses]
        expected_result = [0, 1, 1, 0]

        self.assertEqual(result, expected_result)
Esempio n. 24
0
    def test_commit(self):
        from google.cloud._testing import _Monkey
        from tests.unit._testing import _FakeStub
        from google.cloud.bigtable import row as MUT

        row_key = b'row_key'
        table_name = 'projects/more-stuff'
        column_family_id = u'column_family_id'
        column = b'column'
        client = _Client()
        table = _Table(table_name, client=client)
        row = self._make_one(row_key, table)

        # Create request_pb
        value = b'bytes-value'
        # We will call row.append_cell_value(COLUMN_FAMILY_ID, COLUMN, value).
        request_pb = _ReadModifyWriteRowRequestPB(
            table_name=table_name,
            row_key=row_key,
            rules=[
                _ReadModifyWriteRulePB(
                    family_name=column_family_id,
                    column_qualifier=column,
                    append_value=value,
                ),
            ],
        )

        # Create response_pb
        response_pb = object()

        # Patch the stub used by the API method.
        client._data_stub = stub = _FakeStub(response_pb)

        # Create expected_result.
        row_responses = []
        expected_result = object()

        def mock_parse_rmw_row_response(row_response):
            row_responses.append(row_response)
            return expected_result

        # Perform the method and check the result.
        with _Monkey(MUT, _parse_rmw_row_response=mock_parse_rmw_row_response):
            row.append_cell_value(column_family_id, column, value)
            result = row.commit()

        self.assertEqual(result, expected_result)
        self.assertEqual(stub.method_calls, [(
            'ReadModifyWriteRow',
            (request_pb,),
            {},
        )])
        self.assertEqual(row_responses, [response_pb])
        self.assertEqual(row._rule_pb_list, [])
Esempio n. 25
0
def _create_test_helper(gc_rule=None):
    from google.cloud.bigtable_admin_v2.types import (
        bigtable_table_admin as table_admin_v2_pb2, )
    from tests.unit._testing import _FakeStub
    from google.cloud.bigtable_admin_v2.services.bigtable_table_admin import (
        BigtableTableAdminClient, )

    project_id = "project-id"
    zone = "zone"
    cluster_id = "cluster-id"
    table_id = "table-id"
    column_family_id = "column-family-id"
    table_name = ("projects/" + project_id + "/zones/" + zone + "/clusters/" +
                  cluster_id + "/tables/" + table_id)

    api = mock.create_autospec(BigtableTableAdminClient)

    credentials = _make_credentials()
    client = _make_client(project=project_id,
                          credentials=credentials,
                          admin=True)
    table = _Table(table_name, client=client)
    column_family = _make_column_family(column_family_id,
                                        table,
                                        gc_rule=gc_rule)

    # Create request_pb
    if gc_rule is None:
        column_family_pb = _ColumnFamilyPB()
    else:
        column_family_pb = _ColumnFamilyPB(gc_rule=gc_rule.to_pb())
    request_pb = table_admin_v2_pb2.ModifyColumnFamiliesRequest(
        name=table_name)
    modification = table_admin_v2_pb2.ModifyColumnFamiliesRequest.Modification(
    )
    modification.id = column_family_id
    modification.create = column_family_pb
    request_pb.modifications.append(modification)

    # Create response_pb
    response_pb = _ColumnFamilyPB()

    # Patch the stub used by the API method.
    stub = _FakeStub(response_pb)
    client._table_admin_client = api
    client._table_admin_client.transport.create = stub

    # Create expected_result.
    expected_result = None  # create() has no return value.

    # Perform the method and check the result.
    assert stub.results == (response_pb, )
    result = column_family.create()
    assert result == expected_result
Esempio n. 26
0
    def test_do_mutate_retryable_rows_second_retry(self):
        from google.cloud.bigtable.row import DirectRow
        from google.cloud.bigtable.table import _BigtableRetryableError
        from tests.unit._testing import _FakeStub

        # Setup:
        #   - Mutate 4 rows.
        #   - First try results:
        #       [success, retryable, non-retryable, retryable]
        # Action:
        #   - Second try should re-attempt the 'retryable' rows.
        # Expectation:
        #   - After second try:
        #       [success, success, non-retryable, retryable]
        #   - One of the rows tried second time returns retryable error code,
        #     so expect a raise.
        #   - Exception contains response whose index should be '3' even though
        #     only two rows were retried.

        client = _Client()
        instance = _Instance(self.INSTANCE_NAME, client=client)
        table = self._make_table(self.TABLE_ID, instance)

        row_1 = DirectRow(row_key=b'row_key', table=table)
        row_1.set_cell('cf', b'col', b'value1')
        row_2 = DirectRow(row_key=b'row_key_2', table=table)
        row_2.set_cell('cf', b'col', b'value2')
        row_3 = DirectRow(row_key=b'row_key_3', table=table)
        row_3.set_cell('cf', b'col', b'value3')
        row_4 = DirectRow(row_key=b'row_key_4', table=table)
        row_4.set_cell('cf', b'col', b'value4')

        response = self._make_responses([self.SUCCESS, self.RETRYABLE_1])

        # Patch the stub used by the API method.
        client._data_stub = _FakeStub([response])

        worker = self._make_worker(table._instance._client, table.name,
                                   [row_1, row_2, row_3, row_4])
        worker.responses_statuses = self._make_responses_statuses([
            self.SUCCESS, self.RETRYABLE_1, self.NON_RETRYABLE,
            self.RETRYABLE_2
        ])

        with self.assertRaises(_BigtableRetryableError):
            worker._do_mutate_retryable_rows()

        statuses = worker.responses_statuses
        result = [status.code for status in statuses]
        expected_result = [
            self.SUCCESS, self.SUCCESS, self.NON_RETRYABLE, self.RETRYABLE_1
        ]

        self.assertEqual(result, expected_result)
Esempio n. 27
0
    def _update_test_helper(self, gc_rule=None):
        from tests.unit._testing import _FakeStub
        from google.cloud.bigtable_admin_v2.proto import (
            bigtable_table_admin_pb2 as table_admin_v2_pb2,
        )
        from google.cloud.bigtable_admin_v2.gapic import bigtable_table_admin_client

        project_id = "project-id"
        zone = "zone"
        cluster_id = "cluster-id"
        table_id = "table-id"
        column_family_id = "column-family-id"
        table_name = (
            "projects/"
            + project_id
            + "/zones/"
            + zone
            + "/clusters/"
            + cluster_id
            + "/tables/"
            + table_id
        )

        api = bigtable_table_admin_client.BigtableTableAdminClient(mock.Mock())
        credentials = _make_credentials()
        client = self._make_client(
            project=project_id, credentials=credentials, admin=True
        )
        table = _Table(table_name, client=client)
        column_family = self._make_one(column_family_id, table, gc_rule=gc_rule)

        # Create request_pb
        if gc_rule is None:
            column_family_pb = _ColumnFamilyPB()
        else:
            column_family_pb = _ColumnFamilyPB(gc_rule=gc_rule.to_pb())
        request_pb = table_admin_v2_pb2.ModifyColumnFamiliesRequest(name=table_name)
        request_pb.modifications.add(id=column_family_id, update=column_family_pb)

        # Create response_pb
        response_pb = _ColumnFamilyPB()

        # Patch the stub used by the API method.
        stub = _FakeStub(response_pb)
        client._table_admin_client = api
        client._table_admin_client.transport.update = stub

        # Create expected_result.
        expected_result = None  # update() has no return value.

        # Perform the method and check the result.
        self.assertEqual(stub.results, (response_pb,))
        result = column_family.update()
        self.assertEqual(result, expected_result)
Esempio n. 28
0
    def test_create(self):
        import datetime
        from google.api_core import operation
        from google.longrunning import operations_pb2
        from google.protobuf.any_pb2 import Any
        from google.cloud.bigtable._generated import (
            bigtable_instance_admin_pb2 as messages_v2_pb2)
        from google.cloud._helpers import _datetime_to_pb_timestamp
        from tests.unit._testing import _FakeStub
        from google.cloud.bigtable.cluster import DEFAULT_SERVE_NODES

        NOW = datetime.datetime.utcnow()
        NOW_PB = _datetime_to_pb_timestamp(NOW)
        client = _Client(self.PROJECT)
        instance = self._make_one(self.INSTANCE_ID, client, self.LOCATION_ID,
                                  display_name=self.DISPLAY_NAME)

        # Create response_pb
        metadata = messages_v2_pb2.CreateInstanceMetadata(request_time=NOW_PB)
        type_url = 'type.googleapis.com/%s' % (
            messages_v2_pb2.CreateInstanceMetadata.DESCRIPTOR.full_name,)
        response_pb = operations_pb2.Operation(
            name=self.OP_NAME,
            metadata=Any(
                type_url=type_url,
                value=metadata.SerializeToString(),
            )
        )

        # Patch the stub used by the API method.
        client._instance_stub = stub = _FakeStub(response_pb)

        # Perform the method and check the result.
        result = instance.create()

        self.assertIsInstance(result, operation.Operation)
        self.assertEqual(result.operation.name, self.OP_NAME)
        self.assertIsInstance(result.metadata,
                              messages_v2_pb2.CreateInstanceMetadata)
        self.assertEqual(result.metadata.request_time, NOW_PB)

        self.assertEqual(len(stub.method_calls), 1)
        api_name, args, kwargs = stub.method_calls[0]
        self.assertEqual(api_name, 'CreateInstance')
        request_pb, = args
        self.assertIsInstance(request_pb,
                              messages_v2_pb2.CreateInstanceRequest)
        self.assertEqual(request_pb.parent, 'projects/%s' % (self.PROJECT,))
        self.assertEqual(request_pb.instance_id, self.INSTANCE_ID)
        self.assertEqual(request_pb.instance.display_name, self.DISPLAY_NAME)
        cluster = request_pb.clusters[self.INSTANCE_ID]
        self.assertEqual(cluster.serve_nodes, DEFAULT_SERVE_NODES)
        self.assertEqual(kwargs, {})
Esempio n. 29
0
    def test_read_rows(self):
        from google.cloud._testing import _Monkey
        from tests.unit._testing import _FakeStub
        from google.cloud.bigtable.row_data import PartialRowsData
        from google.cloud.bigtable import table as MUT

        client = _Client()
        instance = _Instance(self.INSTANCE_NAME, client=client)
        table = self._make_one(self.TABLE_ID, instance)

        # Create request_pb
        request_pb = object()  # Returned by our mock.
        mock_created = []

        def mock_create_row_request(table_name, **kwargs):
            mock_created.append((table_name, kwargs))
            return request_pb

        # Create response_iterator
        response_iterator = object()

        # Patch the stub used by the API method.
        client._data_stub = stub = _FakeStub(response_iterator)

        # Create expected_result.
        expected_result = PartialRowsData(response_iterator)

        # Perform the method and check the result.
        start_key = b'start-key'
        end_key = b'end-key'
        filter_obj = object()
        limit = 22
        with _Monkey(MUT, _create_row_request=mock_create_row_request):
            result = table.read_rows(start_key=start_key,
                                     end_key=end_key,
                                     filter_=filter_obj,
                                     limit=limit)

        self.assertEqual(result, expected_result)
        self.assertEqual(stub.method_calls, [(
            'ReadRows',
            (request_pb, ),
            {},
        )])
        created_kwargs = {
            'start_key': start_key,
            'end_key': end_key,
            'filter_': filter_obj,
            'limit': limit,
            'end_inclusive': False,
        }
        self.assertEqual(mock_created, [(table.name, created_kwargs)])
Esempio n. 30
0
    def test_update_app_profile_with_single_routing_policy(self):
        import datetime
        from google.api_core import operation
        from google.longrunning import operations_pb2
        from google.protobuf.any_pb2 import Any
        from google.cloud.bigtable_admin_v2.proto import (
            bigtable_instance_admin_pb2 as messages_v2_pb2)
        from google.cloud._helpers import _datetime_to_pb_timestamp
        from tests.unit._testing import _FakeStub
        from google.cloud.bigtable_admin_v2.gapic import (
            bigtable_instance_admin_client)

        instance_api = (
            bigtable_instance_admin_client.BigtableInstanceAdminClient(
                mock.Mock()))

        credentials = _make_credentials()
        client = self._make_client(project=self.PROJECT,
                                   credentials=credentials,
                                   admin=True)
        instance = self._make_one(self.INSTANCE_ID, client)

        # Create response_pb
        NOW = datetime.datetime.utcnow()
        NOW_PB = _datetime_to_pb_timestamp(NOW)
        metadata = messages_v2_pb2.CreateInstanceMetadata(request_time=NOW_PB)
        type_url = 'type.googleapis.com/%s' % (
            messages_v2_pb2.CreateInstanceMetadata.DESCRIPTOR.full_name, )
        response_pb = operations_pb2.Operation(
            name=self.OP_NAME,
            metadata=Any(
                type_url=type_url,
                value=metadata.SerializeToString(),
            ))

        # Patch the stub used by the API method.
        client._instance_admin_client = instance_api
        stub = _FakeStub(response_pb)
        client._instance_admin_client.bigtable_instance_admin_stub = stub
        update_mask = []

        # Perform the method and check the result.
        app_profile_id = 'appProfileId1262094415'
        cluster_id = 'cluster-id'
        result = instance.update_app_profile(app_profile_id,
                                             update_mask=update_mask,
                                             routing_policy_type=2,
                                             cluster_id=cluster_id)

        self.assertIsInstance(result, operation.Operation)
Esempio n. 31
0
    def test_do_mutate_retryable_rows_second_try(self):
        from google.cloud.bigtable.row import DirectRow
        from tests.unit._testing import _FakeStub

        # Setup:
        #   - Mutate 4 rows.
        #   - First try results:
        #       [success, retryable, non-retryable, retryable]
        # Action:
        #   - Second try should re-attempt the 'retryable' rows.
        # Expectation:
        #   - After second try:
        #       [success, non-retryable, non-retryable, success]

        client = _Client()
        instance = _Instance(self.INSTANCE_NAME, client=client)
        table = self._make_table(self.TABLE_ID, instance)

        row_1 = DirectRow(row_key=b'row_key', table=table)
        row_1.set_cell('cf', b'col', b'value1')
        row_2 = DirectRow(row_key=b'row_key_2', table=table)
        row_2.set_cell('cf', b'col', b'value2')
        row_3 = DirectRow(row_key=b'row_key_3', table=table)
        row_3.set_cell('cf', b'col', b'value3')
        row_4 = DirectRow(row_key=b'row_key_4', table=table)
        row_4.set_cell('cf', b'col', b'value4')

        response = self._make_responses([self.NON_RETRYABLE, self.SUCCESS])

        # Patch the stub used by the API method.
        client._data_stub = _FakeStub([response])

        worker = self._make_worker(
            table._instance._client,
            table.name, [row_1, row_2, row_3, row_4])
        worker.responses_statuses = self._make_responses_statuses([
            self.SUCCESS,
            self.RETRYABLE_1,
            self.NON_RETRYABLE,
            self.RETRYABLE_2])

        statuses = worker._do_mutate_retryable_rows()

        result = [status.code for status in statuses]
        expected_result = [self.SUCCESS,
                           self.NON_RETRYABLE,
                           self.NON_RETRYABLE,
                           self.SUCCESS]

        self.assertEqual(result, expected_result)
    def _update_test_helper(self, gc_rule=None):
        from tests.unit._testing import _FakeStub
        from google.cloud.bigtable._generated import (bigtable_table_admin_pb2
                                                      as table_admin_v2_pb2)

        project_id = 'project-id'
        zone = 'zone'
        cluster_id = 'cluster-id'
        table_id = 'table-id'
        column_family_id = 'column-family-id'
        table_name = ('projects/' + project_id + '/zones/' + zone +
                      '/clusters/' + cluster_id + '/tables/' + table_id)

        client = _Client()
        table = _Table(table_name, client=client)
        column_family = self._make_one(column_family_id,
                                       table,
                                       gc_rule=gc_rule)

        # Create request_pb
        if gc_rule is None:
            column_family_pb = _ColumnFamilyPB()
        else:
            column_family_pb = _ColumnFamilyPB(gc_rule=gc_rule.to_pb())
        request_pb = table_admin_v2_pb2.ModifyColumnFamiliesRequest(
            name=table_name)
        request_pb.modifications.add(
            id=column_family_id,
            update=column_family_pb,
        )

        # Create response_pb
        response_pb = _ColumnFamilyPB()

        # Patch the stub used by the API method.
        client._table_stub = stub = _FakeStub(response_pb)

        # Create expected_result.
        expected_result = None  # update() has no return value.

        # Perform the method and check the result.
        self.assertEqual(stub.results, (response_pb, ))
        result = column_family.update()
        self.assertEqual(stub.results, ())
        self.assertEqual(result, expected_result)
        self.assertEqual(stub.method_calls, [(
            'ModifyColumnFamilies',
            (request_pb, ),
            {},
        )])
    def _update_test_helper(self, gc_rule=None):
        from tests.unit._testing import _FakeStub
        from google.cloud.bigtable._generated import (
            bigtable_table_admin_pb2 as table_admin_v2_pb2)

        project_id = 'project-id'
        zone = 'zone'
        cluster_id = 'cluster-id'
        table_id = 'table-id'
        column_family_id = 'column-family-id'
        table_name = ('projects/' + project_id + '/zones/' + zone +
                      '/clusters/' + cluster_id + '/tables/' + table_id)

        client = _Client()
        table = _Table(table_name, client=client)
        column_family = self._make_one(
            column_family_id, table, gc_rule=gc_rule)

        # Create request_pb
        if gc_rule is None:
            column_family_pb = _ColumnFamilyPB()
        else:
            column_family_pb = _ColumnFamilyPB(gc_rule=gc_rule.to_pb())
        request_pb = table_admin_v2_pb2.ModifyColumnFamiliesRequest(
            name=table_name)
        request_pb.modifications.add(
            id=column_family_id,
            update=column_family_pb,
        )

        # Create response_pb
        response_pb = _ColumnFamilyPB()

        # Patch the stub used by the API method.
        client._table_stub = stub = _FakeStub(response_pb)

        # Create expected_result.
        expected_result = None  # update() has no return value.

        # Perform the method and check the result.
        self.assertEqual(stub.results, (response_pb,))
        result = column_family.update()
        self.assertEqual(stub.results, ())
        self.assertEqual(result, expected_result)
        self.assertEqual(stub.method_calls, [(
            'ModifyColumnFamilies',
            (request_pb,),
            {},
        )])
Esempio n. 34
0
    def test_create(self):
        import datetime
        from google.api_core import operation
        from google.longrunning import operations_pb2
        from google.protobuf.any_pb2 import Any
        from google.cloud.bigtable_admin_v2.proto import (
            bigtable_instance_admin_pb2 as messages_v2_pb2)
        from google.cloud._helpers import _datetime_to_pb_timestamp
        from tests.unit._testing import _FakeStub
        from google.cloud.bigtable_admin_v2.gapic import (
            bigtable_instance_admin_client)

        NOW = datetime.datetime.utcnow()
        NOW_PB = _datetime_to_pb_timestamp(NOW)
        api = bigtable_instance_admin_client.BigtableInstanceAdminClient(
            mock.Mock())
        credentials = _make_credentials()
        client = self._make_client(project=self.PROJECT,
                                   credentials=credentials,
                                   admin=True)
        instance = self._make_one(self.INSTANCE_ID,
                                  client,
                                  self.LOCATION_ID,
                                  display_name=self.DISPLAY_NAME)

        # Create response_pb
        metadata = messages_v2_pb2.CreateInstanceMetadata(request_time=NOW_PB)
        type_url = 'type.googleapis.com/%s' % (
            messages_v2_pb2.CreateInstanceMetadata.DESCRIPTOR.full_name, )
        response_pb = operations_pb2.Operation(
            name=self.OP_NAME,
            metadata=Any(
                type_url=type_url,
                value=metadata.SerializeToString(),
            ))

        # Patch the stub used by the API method.
        stub = _FakeStub(response_pb)
        client._instance_admin_client = api
        client._instance_admin_client.bigtable_instance_admin_stub = stub

        # Perform the method and check the result.
        result = instance.create()

        self.assertIsInstance(result, operation.Operation)
        # self.assertEqual(result.operation.name, self.OP_NAME)
        self.assertIsInstance(result.metadata,
                              messages_v2_pb2.CreateInstanceMetadata)
Esempio n. 35
0
    def test_commit(self):
        from google.protobuf import empty_pb2
        from tests.unit._testing import _FakeStub

        row_key = b'row_key'
        table_name = 'projects/more-stuff'
        column_family_id = u'column_family_id'
        column = b'column'
        client = _Client()
        table = _Table(table_name, client=client)
        row = self._make_one(row_key, table)

        # Create request_pb
        value = b'bytes-value'
        mutation = _MutationPB(
            set_cell=_MutationSetCellPB(
                family_name=column_family_id,
                column_qualifier=column,
                timestamp_micros=-1,  # Default value.
                value=value,
            ),
        )
        request_pb = _MutateRowRequestPB(
            table_name=table_name,
            row_key=row_key,
            mutations=[mutation],
        )

        # Create response_pb
        response_pb = empty_pb2.Empty()

        # Patch the stub used by the API method.
        client._data_stub = stub = _FakeStub(response_pb)

        # Create expected_result.
        expected_result = None  # commit() has no return value when no filter.

        # Perform the method and check the result.
        row.set_cell(column_family_id, column, value)
        result = row.commit()
        self.assertEqual(result, expected_result)
        self.assertEqual(stub.method_calls, [(
            'MutateRow',
            (request_pb,),
            {},
        )])
        self.assertEqual(row._pb_mutations, [])
Esempio n. 36
0
    def test_do_mutate_retryable_rows_second_try(self):
        from google.cloud.bigtable.row import DirectRow
        from tests.unit._testing import _FakeStub

        # Setup:
        #   - Mutate 4 rows.
        #   - First try results:
        #       [success, retryable, non-retryable, retryable]
        # Action:
        #   - Second try should re-attempt the 'retryable' rows.
        # Expectation:
        #   - After second try:
        #       [success, non-retryable, non-retryable, success]

        client = _Client()
        instance = _Instance(self.INSTANCE_NAME, client=client)
        table = self._make_table(self.TABLE_ID, instance)

        row_1 = DirectRow(row_key=b'row_key', table=table)
        row_1.set_cell('cf', b'col', b'value1')
        row_2 = DirectRow(row_key=b'row_key_2', table=table)
        row_2.set_cell('cf', b'col', b'value2')
        row_3 = DirectRow(row_key=b'row_key_3', table=table)
        row_3.set_cell('cf', b'col', b'value3')
        row_4 = DirectRow(row_key=b'row_key_4', table=table)
        row_4.set_cell('cf', b'col', b'value4')

        response = self._make_responses([self.NON_RETRYABLE, self.SUCCESS])

        # Patch the stub used by the API method.
        client._data_stub = _FakeStub([response])

        worker = self._make_worker(table._instance._client, table.name,
                                   [row_1, row_2, row_3, row_4])
        worker.responses_statuses = self._make_responses_statuses([
            self.SUCCESS, self.RETRYABLE_1, self.NON_RETRYABLE,
            self.RETRYABLE_2
        ])

        statuses = worker._do_mutate_retryable_rows()

        result = [status.code for status in statuses]
        expected_result = [
            self.SUCCESS, self.NON_RETRYABLE, self.NON_RETRYABLE, self.SUCCESS
        ]

        self.assertEqual(result, expected_result)
Esempio n. 37
0
    def test_read_rows_one_chunk(self):
        from google.cloud._testing import _Monkey
        from tests.unit._testing import _FakeStub
        from google.cloud.bigtable import retry as MUT
        from google.cloud.bigtable.retry import ReadRowsIterator
        from google.cloud.bigtable.row_data import Cell
        from google.cloud.bigtable.row_data import PartialRowsData

        client = _Client()
        instance = _Instance(self.INSTANCE_NAME, client=client)
        table = self._make_one(self.TABLE_ID, instance)

        # Create request_pb
        request_pb = object()  # Returned by our mock.
        mock_created = []

        def mock_create_row_request(table_name, **kwargs):
            mock_created.append((table_name, kwargs))
            return request_pb

        # Create response_iterator
        chunk = _ReadRowsResponseCellChunkPB(
            row_key=self.ROW_KEY,
            family_name=self.FAMILY_NAME,
            qualifier=self.QUALIFIER,
            timestamp_micros=self.TIMESTAMP_MICROS,
            value=self.VALUE,
            commit_row=True,
        )
        response_pb = _ReadRowsResponsePB(chunks=[chunk])
        response_iterator = iter([response_pb])

        # Patch the stub used by the API method.
        client._data_stub = stub = _FakeStub(response_iterator)

        start_key = b'start-key'
        end_key = b'end-key'
        filter_obj = object()
        limit = 22
        with _Monkey(MUT, _create_row_request=mock_create_row_request):
            # Perform the method and check the result.
            result = table.read_rows(start_key=start_key,
                                     end_key=end_key,
                                     filter_=filter_obj,
                                     limit=limit)
            result.consume_all()
Esempio n. 38
0
    def test_commit_no_rules(self):
        from tests.unit._testing import _FakeStub

        row_key = b'row_key'
        client = _Client()
        table = _Table(None, client=client)
        row = self._make_one(row_key, table)
        self.assertEqual(row._rule_pb_list, [])

        # Patch the stub used by the API method.
        client._data_stub = stub = _FakeStub()

        # Perform the method and check the result.
        result = row.commit()
        self.assertEqual(result, {})
        # Make sure no request was sent.
        self.assertEqual(stub.method_calls, [])
Esempio n. 39
0
    def test_delete(self):
        from google.protobuf import empty_pb2
        from google.cloud.bigtable_admin_v2.types import (
            bigtable_table_admin as table_admin_v2_pb2, )
        from tests.unit._testing import _FakeStub
        from google.cloud.bigtable_admin_v2.services.bigtable_table_admin import (
            BigtableTableAdminClient, )

        project_id = "project-id"
        zone = "zone"
        cluster_id = "cluster-id"
        table_id = "table-id"
        column_family_id = "column-family-id"
        table_name = ("projects/" + project_id + "/zones/" + zone +
                      "/clusters/" + cluster_id + "/tables/" + table_id)

        api = mock.create_autospec(BigtableTableAdminClient)
        credentials = _make_credentials()
        client = self._make_client(project=project_id,
                                   credentials=credentials,
                                   admin=True)
        table = _Table(table_name, client=client)
        column_family = self._make_one(column_family_id, table)

        # Create request_pb
        request_pb = table_admin_v2_pb2.ModifyColumnFamiliesRequest(
            name=table_name)
        modification = table_admin_v2_pb2.ModifyColumnFamiliesRequest.Modification(
            id=column_family_id, drop=True)
        request_pb.modifications.append(modification)

        # Create response_pb
        response_pb = empty_pb2.Empty()

        # Patch the stub used by the API method.
        stub = _FakeStub(response_pb)
        client._table_admin_client = api
        client._table_admin_client.transport.delete = stub

        # Create expected_result.
        expected_result = None  # delete() has no return value.

        # Perform the method and check the result.
        self.assertEqual(stub.results, (response_pb, ))
        result = column_family.delete()
        self.assertEqual(result, expected_result)
Esempio n. 40
0
    def test_commit_no_rules(self):
        from tests.unit._testing import _FakeStub

        row_key = b'row_key'
        client = _Client()
        table = _Table(None, client=client)
        row = self._make_one(row_key, table)
        self.assertEqual(row._rule_pb_list, [])

        # Patch the stub used by the API method.
        client._data_stub = stub = _FakeStub()

        # Perform the method and check the result.
        result = row.commit()
        self.assertEqual(result, {})
        # Make sure no request was sent.
        self.assertEqual(stub.method_calls, [])
    def test_delete(self):
        from google.protobuf import empty_pb2
        from google.cloud.bigtable_admin_v2.proto import (
            bigtable_table_admin_pb2 as table_admin_v2_pb2)
        from tests.unit._testing import _FakeStub
        from google.cloud.bigtable_admin_v2.gapic import (
            bigtable_table_admin_client)

        project_id = 'project-id'
        zone = 'zone'
        cluster_id = 'cluster-id'
        table_id = 'table-id'
        column_family_id = 'column-family-id'
        table_name = ('projects/' + project_id + '/zones/' + zone +
                      '/clusters/' + cluster_id + '/tables/' + table_id)

        api = bigtable_table_admin_client.BigtableTableAdminClient(mock.Mock())
        credentials = _make_credentials()
        client = self._make_client(project=project_id,
                                   credentials=credentials, admin=True)
        table = _Table(table_name, client=client)
        column_family = self._make_one(column_family_id, table)

        # Create request_pb
        request_pb = table_admin_v2_pb2.ModifyColumnFamiliesRequest(
            name=table_name)
        request_pb.modifications.add(
            id=column_family_id,
            drop=True)

        # Create response_pb
        response_pb = empty_pb2.Empty()

        # Patch the stub used by the API method.
        stub = _FakeStub(response_pb)
        client._table_admin_client = api
        client._table_admin_client.bigtable_table_admin_stub = stub

        # Create expected_result.
        expected_result = None  # delete() has no return value.

        # Perform the method and check the result.
        self.assertEqual(stub.results, (response_pb,))
        result = column_family.delete()
        self.assertEqual(result, expected_result)
Esempio n. 42
0
    def test_do_mutate_retryable_rows(self):
        from google.cloud.bigtable._generated.bigtable_pb2 import (
            MutateRowsResponse)
        from google.cloud.bigtable.row import DirectRow
        from google.rpc.status_pb2 import Status
        from tests.unit._testing import _FakeStub

        # Setup:
        #   - Mutate 2 rows.
        # Action:
        #   - Initial attempt will mutate all 2 rows.
        # Expectation:
        #   - Expect [success, non-retryable]

        client = _Client()
        instance = _Instance(self.INSTANCE_NAME, client=client)
        table = self._make_table(self.TABLE_ID, instance)

        row_1 = DirectRow(row_key=b'row_key', table=table)
        row_1.set_cell('cf', b'col', b'value1')
        row_2 = DirectRow(row_key=b'row_key_2', table=table)
        row_2.set_cell('cf', b'col', b'value2')

        response = MutateRowsResponse(entries=[
            MutateRowsResponse.Entry(
                index=0,
                status=Status(code=0),
            ),
            MutateRowsResponse.Entry(
                index=1,
                status=Status(code=1),
            ),
        ], )

        # Patch the stub used by the API method.
        client._data_stub = _FakeStub([response])

        worker = self._make_worker(table._instance._client, table.name,
                                   [row_1, row_2])
        statuses = worker._do_mutate_retryable_rows()

        result = [status.code for status in statuses]
        expected_result = [0, 1]

        self.assertEqual(result, expected_result)
Esempio n. 43
0
    def test_delete(self):
        from google.protobuf import empty_pb2
        from google.cloud.bigtable_admin_v2.proto import (
            bigtable_table_admin_pb2 as table_admin_v2_pb2)
        from tests.unit._testing import _FakeStub
        from google.cloud.bigtable_admin_v2.gapic import (
            bigtable_table_admin_client)

        project_id = 'project-id'
        zone = 'zone'
        cluster_id = 'cluster-id'
        table_id = 'table-id'
        column_family_id = 'column-family-id'
        table_name = ('projects/' + project_id + '/zones/' + zone +
                      '/clusters/' + cluster_id + '/tables/' + table_id)

        api = bigtable_table_admin_client.BigtableTableAdminClient(mock.Mock())
        credentials = _make_credentials()
        client = self._make_client(project=project_id,
                                   credentials=credentials, admin=True)
        table = _Table(table_name, client=client)
        column_family = self._make_one(column_family_id, table)

        # Create request_pb
        request_pb = table_admin_v2_pb2.ModifyColumnFamiliesRequest(
            name=table_name)
        request_pb.modifications.add(
            id=column_family_id,
            drop=True)

        # Create response_pb
        response_pb = empty_pb2.Empty()

        # Patch the stub used by the API method.
        stub = _FakeStub(response_pb)
        client._table_admin_client = api
        client._table_admin_client.transport.delete = stub

        # Create expected_result.
        expected_result = None  # delete() has no return value.

        # Perform the method and check the result.
        self.assertEqual(stub.results, (response_pb,))
        result = column_family.delete()
        self.assertEqual(result, expected_result)
Esempio n. 44
0
    def _create_test_helper(self, initial_split_keys, column_families=()):
        from google.cloud._helpers import _to_bytes
        from tests.unit._testing import _FakeStub

        client = _Client()
        instance = _Instance(self.INSTANCE_NAME, client=client)
        table = self._make_one(self.TABLE_ID, instance)

        # Create request_pb
        splits_pb = [
            _CreateTableRequestSplitPB(key=_to_bytes(key))
            for key in initial_split_keys or ()
        ]
        table_pb = None
        if column_families:
            table_pb = _TablePB()
            for cf in column_families:
                cf_pb = table_pb.column_families[cf.column_family_id]
                if cf.gc_rule is not None:
                    cf_pb.gc_rule.MergeFrom(cf.gc_rule.to_pb())
        request_pb = _CreateTableRequestPB(
            initial_splits=splits_pb,
            parent=self.INSTANCE_NAME,
            table_id=self.TABLE_ID,
            table=table_pb,
        )

        # Create response_pb
        response_pb = _TablePB()

        # Patch the stub used by the API method.
        client._table_stub = stub = _FakeStub(response_pb)

        # Create expected_result.
        expected_result = None  # create() has no return value.

        # Perform the method and check the result.
        result = table.create(initial_split_keys=initial_split_keys,
                              column_families=column_families)
        self.assertEqual(result, expected_result)
        self.assertEqual(stub.method_calls, [(
            'CreateTable',
            (request_pb, ),
            {},
        )])
Esempio n. 45
0
    def _create_test_helper(self, initial_split_keys, column_families=()):
        from google.cloud._helpers import _to_bytes
        from tests.unit._testing import _FakeStub

        client = _Client()
        instance = _Instance(self.INSTANCE_NAME, client=client)
        table = self._make_one(self.TABLE_ID, instance)

        # Create request_pb
        splits_pb = [
            _CreateTableRequestSplitPB(key=_to_bytes(key))
            for key in initial_split_keys or ()]
        table_pb = None
        if column_families:
            table_pb = _TablePB()
            for cf in column_families:
                cf_pb = table_pb.column_families[cf.column_family_id]
                if cf.gc_rule is not None:
                    cf_pb.gc_rule.CopyFrom(cf.gc_rule.to_pb())
        request_pb = _CreateTableRequestPB(
            initial_splits=splits_pb,
            parent=self.INSTANCE_NAME,
            table_id=self.TABLE_ID,
            table=table_pb,
        )

        # Create response_pb
        response_pb = _TablePB()

        # Patch the stub used by the API method.
        client._table_stub = stub = _FakeStub(response_pb)

        # Create expected_result.
        expected_result = None  # create() has no return value.

        # Perform the method and check the result.
        result = table.create(initial_split_keys=initial_split_keys,
                              column_families=column_families)
        self.assertEqual(result, expected_result)
        self.assertEqual(stub.method_calls, [(
            'CreateTable',
            (request_pb,),
            {},
        )])
Esempio n. 46
0
    def test_do_mutate_retryable_rows_retry(self):
        from google.cloud.bigtable.row import DirectRow
        from google.cloud.bigtable.table import _BigtableRetryableError
        from tests.unit._testing import _FakeStub

        # Setup:
        #   - Mutate 3 rows.
        # Action:
        #   - Initial attempt will mutate all 3 rows.
        # Expectation:
        #   - Second row returns retryable error code, so expect a raise.
        #   - State of responses_statuses should be
        #       [success, retryable, non-retryable]

        client = _Client()
        instance = _Instance(self.INSTANCE_NAME, client=client)
        table = self._make_table(self.TABLE_ID, instance)

        row_1 = DirectRow(row_key=b'row_key', table=table)
        row_1.set_cell('cf', b'col', b'value1')
        row_2 = DirectRow(row_key=b'row_key_2', table=table)
        row_2.set_cell('cf', b'col', b'value2')
        row_3 = DirectRow(row_key=b'row_key_3', table=table)
        row_3.set_cell('cf', b'col', b'value3')

        response = self._make_responses([
            self.SUCCESS,
            self.RETRYABLE_1,
            self.NON_RETRYABLE])

        # Patch the stub used by the API method.
        client._data_stub = _FakeStub([response])

        worker = self._make_worker(
            table._instance._client, table.name, [row_1, row_2, row_3])

        with self.assertRaises(_BigtableRetryableError):
            worker._do_mutate_retryable_rows()

        statuses = worker.responses_statuses
        result = [status.code for status in statuses]
        expected_result = [self.SUCCESS, self.RETRYABLE_1, self.NON_RETRYABLE]

        self.assertEqual(result, expected_result)
    def test_delete(self):
        from google.protobuf import empty_pb2
        from google.cloud.bigtable._generated import (
            bigtable_table_admin_pb2 as table_admin_v2_pb2)
        from tests.unit._testing import _FakeStub

        project_id = 'project-id'
        zone = 'zone'
        cluster_id = 'cluster-id'
        table_id = 'table-id'
        column_family_id = 'column-family-id'
        table_name = ('projects/' + project_id + '/zones/' + zone +
                      '/clusters/' + cluster_id + '/tables/' + table_id)

        client = _Client()
        table = _Table(table_name, client=client)
        column_family = self._make_one(column_family_id, table)

        # Create request_pb
        request_pb = table_admin_v2_pb2.ModifyColumnFamiliesRequest(
            name=table_name)
        request_pb.modifications.add(
            id=column_family_id,
            drop=True)

        # Create response_pb
        response_pb = empty_pb2.Empty()

        # Patch the stub used by the API method.
        client._table_stub = stub = _FakeStub(response_pb)

        # Create expected_result.
        expected_result = None  # delete() has no return value.

        # Perform the method and check the result.
        self.assertEqual(stub.results, (response_pb,))
        result = column_family.delete()
        self.assertEqual(stub.results, ())
        self.assertEqual(result, expected_result)
        self.assertEqual(stub.method_calls, [(
            'ModifyColumnFamilies',
            (request_pb,),
            {},
        )])
Esempio n. 48
0
    def test_commit_no_mutations(self):
        from tests.unit._testing import _FakeStub

        row_key = b'row_key'
        client = _Client()
        table = _Table(None, client=client)
        filter_ = object()
        row = self._make_one(row_key, table, filter_=filter_)
        self.assertEqual(row._true_pb_mutations, [])
        self.assertEqual(row._false_pb_mutations, [])

        # Patch the stub used by the API method.
        client._data_stub = stub = _FakeStub()

        # Perform the method and check the result.
        result = row.commit()
        self.assertIsNone(result)
        # Make sure no request was sent.
        self.assertEqual(stub.method_calls, [])
Esempio n. 49
0
    def test_commit_no_rules(self):
        from tests.unit._testing import _FakeStub

        project_id = 'project-id'
        row_key = b'row_key'

        credentials = _make_credentials()
        client = self._make_client(project=project_id, credentials=credentials,
                                   admin=True)
        table = _Table(None, client=client)
        row = self._make_one(row_key, table)
        self.assertEqual(row._rule_pb_list, [])

        # Patch the stub used by the API method.
        stub = _FakeStub()

        # Perform the method and check the result.
        result = row.commit()
        self.assertEqual(result, {})
        # Make sure no request was sent.
        self.assertEqual(stub.method_calls, [])
Esempio n. 50
0
    def test_reload(self):
        from google.cloud.bigtable._generated import (
            instance_pb2 as data_v2_pb2)
        from google.cloud.bigtable._generated import (
            bigtable_instance_admin_pb2 as messages_v2_pb)
        from tests.unit._testing import _FakeStub

        client = _Client(self.PROJECT)
        instance = self._make_one(self.INSTANCE_ID, client, self.LOCATION_ID)

        # Create request_pb
        request_pb = messages_v2_pb.GetInstanceRequest(
            name=self.INSTANCE_NAME)

        # Create response_pb
        DISPLAY_NAME = u'hey-hi-hello'
        response_pb = data_v2_pb2.Instance(
            display_name=DISPLAY_NAME,
        )

        # Patch the stub used by the API method.
        client._instance_stub = stub = _FakeStub(response_pb)

        # Create expected_result.
        expected_result = None  # reload() has no return value.

        # Check Instance optional config values before.
        self.assertEqual(instance.display_name, self.INSTANCE_ID)

        # Perform the method and check the result.
        result = instance.reload()
        self.assertEqual(result, expected_result)
        self.assertEqual(stub.method_calls, [(
            'GetInstance',
            (request_pb,),
            {},
        )])

        # Check Instance optional config values before.
        self.assertEqual(instance.display_name, DISPLAY_NAME)
Esempio n. 51
0
    def _list_tables_helper(self, table_name=None):
        from google.cloud.bigtable._generated import (
            table_pb2 as table_data_v2_pb2)
        from google.cloud.bigtable._generated import (
            bigtable_table_admin_pb2 as table_messages_v1_pb2)
        from tests.unit._testing import _FakeStub

        client = _Client(self.PROJECT)
        instance = self._make_one(self.INSTANCE_ID, client, self.LOCATION_ID)

        # Create request_
        request_pb = table_messages_v1_pb2.ListTablesRequest(
            parent=self.INSTANCE_NAME)

        # Create response_pb
        if table_name is None:
            table_name = self.TABLE_NAME

        response_pb = table_messages_v1_pb2.ListTablesResponse(
            tables=[
                table_data_v2_pb2.Table(name=table_name),
            ],
        )

        # Patch the stub used by the API method.
        client._table_stub = stub = _FakeStub(response_pb)

        # Create expected_result.
        expected_table = instance.table(self.TABLE_ID)
        expected_result = [expected_table]

        # Perform the method and check the result.
        result = instance.list_tables()

        self.assertEqual(result, expected_result)
        self.assertEqual(stub.method_calls, [(
            'ListTables',
            (request_pb,),
            {},
        )])
Esempio n. 52
0
    def _read_row_helper(self, chunks, expected_result):
        from google.cloud._testing import _Monkey
        from tests.unit._testing import _FakeStub
        from google.cloud.bigtable import table as MUT

        client = _Client()
        instance = _Instance(self.INSTANCE_NAME, client=client)
        table = self._make_one(self.TABLE_ID, instance)

        # Create request_pb
        request_pb = object()  # Returned by our mock.
        mock_created = []

        def mock_create_row_request(table_name, row_key, filter_):
            mock_created.append((table_name, row_key, filter_))
            return request_pb

        # Create response_iterator
        if chunks is None:
            response_iterator = iter(())  # no responses at all
        else:
            response_pb = _ReadRowsResponsePB(chunks=chunks)
            response_iterator = iter([response_pb])

        # Patch the stub used by the API method.
        client._data_stub = stub = _FakeStub(response_iterator)

        # Perform the method and check the result.
        filter_obj = object()
        with _Monkey(MUT, _create_row_request=mock_create_row_request):
            result = table.read_row(self.ROW_KEY, filter_=filter_obj)

        self.assertEqual(result, expected_result)
        self.assertEqual(stub.method_calls, [(
            'ReadRows',
            (request_pb,),
            {},
        )])
        self.assertEqual(mock_created,
                         [(table.name, self.ROW_KEY, filter_obj)])
Esempio n. 53
0
    def test_create(self):
        from google.api_core import operation
        from google.longrunning import operations_pb2
        from google.cloud.bigtable._generated import (
            bigtable_instance_admin_pb2 as messages_v2_pb2)
        from tests.unit._testing import _FakeStub

        SERVE_NODES = 4
        client = _Client(self.PROJECT)
        instance = _Instance(self.INSTANCE_ID, client)
        cluster = self._make_one(
            self.CLUSTER_ID, instance, serve_nodes=SERVE_NODES)

        # Create response_pb
        OP_ID = 5678
        OP_NAME = (
            'operations/projects/%s/instances/%s/clusters/%s/operations/%d' %
            (self.PROJECT, self.INSTANCE_ID, self.CLUSTER_ID, OP_ID))
        response_pb = operations_pb2.Operation(name=OP_NAME)

        # Patch the stub used by the API method.
        client._instance_stub = stub = _FakeStub(response_pb)

        # Perform the method and check the result.
        result = cluster.create()

        self.assertIsInstance(result, operation.Operation)
        self.assertEqual(result.operation.name, OP_NAME)
        self.assertIsNone(result.metadata)

        self.assertEqual(len(stub.method_calls), 1)
        api_name, args, kwargs = stub.method_calls[0]
        self.assertEqual(api_name, 'CreateCluster')
        request_pb, = args
        self.assertIsInstance(request_pb,
                              messages_v2_pb2.CreateClusterRequest)
        self.assertEqual(request_pb.parent, instance.name)
        self.assertEqual(request_pb.cluster_id, self.CLUSTER_ID)
        self.assertEqual(request_pb.cluster.serve_nodes, SERVE_NODES)
        self.assertEqual(kwargs, {})
Esempio n. 54
0
    def test_reload(self):
        from tests.unit._testing import _FakeStub
        from google.cloud.bigtable.cluster import DEFAULT_SERVE_NODES

        SERVE_NODES = 31
        LOCATION = 'LOCATION'
        client = _Client(self.PROJECT)
        instance = _Instance(self.INSTANCE_ID, client)
        cluster = self._make_one(self.CLUSTER_ID, instance)

        # Create request_pb
        request_pb = _GetClusterRequestPB(name=self.CLUSTER_NAME)

        # Create response_pb
        response_pb = _ClusterPB(
            serve_nodes=SERVE_NODES,
            location=LOCATION,
        )

        # Patch the stub used by the API method.
        client._instance_stub = stub = _FakeStub(response_pb)

        # Create expected_result.
        expected_result = None  # reload() has no return value.

        # Check Cluster optional config values before.
        self.assertEqual(cluster.serve_nodes, DEFAULT_SERVE_NODES)

        # Perform the method and check the result.
        result = cluster.reload()
        self.assertEqual(result, expected_result)
        self.assertEqual(stub.method_calls, [(
            'GetCluster',
            (request_pb,),
            {},
        )])

        # Check Cluster optional config values before.
        self.assertEqual(cluster.serve_nodes, SERVE_NODES)
        self.assertEqual(cluster.location, LOCATION)
Esempio n. 55
0
    def test_do_mutate_retryable_rows_mismatch_num_responses(self):
        from google.cloud.bigtable.row import DirectRow
        from tests.unit._testing import _FakeStub

        client = _Client()
        instance = _Instance(self.INSTANCE_NAME, client=client)
        table = self._make_table(self.TABLE_ID, instance)

        row_1 = DirectRow(row_key=b'row_key', table=table)
        row_1.set_cell('cf', b'col', b'value1')
        row_2 = DirectRow(row_key=b'row_key_2', table=table)
        row_2.set_cell('cf', b'col', b'value2')

        response = self._make_responses([self.SUCCESS])

        # Patch the stub used by the API method.
        client._data_stub = _FakeStub([response])

        worker = self._make_worker(
            table._instance._client,
            table.name, [row_1, row_2])
        with self.assertRaises(RuntimeError):
            worker._do_mutate_retryable_rows()
Esempio n. 56
0
    def test_create_w_explicit_serve_nodes(self):
        from google.api_core import operation
        from google.longrunning import operations_pb2
        from google.cloud.bigtable._generated import (
            bigtable_instance_admin_pb2 as messages_v2_pb2)
        from tests.unit._testing import _FakeStub

        SERVE_NODES = 5

        client = _Client(self.PROJECT)
        instance = self._make_one(self.INSTANCE_ID, client, self.LOCATION_ID,
                                  serve_nodes=SERVE_NODES)

        # Create response_pb
        response_pb = operations_pb2.Operation(name=self.OP_NAME)

        # Patch the stub used by the API method.
        client._instance_stub = stub = _FakeStub(response_pb)

        # Perform the method and check the result.
        result = instance.create()

        self.assertIsInstance(result, operation.Operation)
        self.assertEqual(result.operation.name, self.OP_NAME)

        self.assertEqual(len(stub.method_calls), 1)
        api_name, args, kwargs = stub.method_calls[0]
        self.assertEqual(api_name, 'CreateInstance')
        request_pb, = args
        self.assertIsInstance(request_pb,
                              messages_v2_pb2.CreateInstanceRequest)
        self.assertEqual(request_pb.parent, 'projects/%s' % (self.PROJECT,))
        self.assertEqual(request_pb.instance_id, self.INSTANCE_ID)
        self.assertEqual(request_pb.instance.display_name, self.INSTANCE_ID)
        cluster = request_pb.clusters[self.INSTANCE_ID]
        self.assertEqual(cluster.serve_nodes, SERVE_NODES)
        self.assertEqual(kwargs, {})
Esempio n. 57
0
    def test_update(self):
        import datetime
        from google.api_core import operation
        from google.longrunning import operations_pb2
        from google.protobuf.any_pb2 import Any
        from google.cloud._helpers import _datetime_to_pb_timestamp
        from google.cloud.bigtable._generated import (
            instance_pb2 as data_v2_pb2)
        from google.cloud.bigtable._generated import (
            bigtable_instance_admin_pb2 as messages_v2_pb2)
        from tests.unit._testing import _FakeStub

        NOW = datetime.datetime.utcnow()
        NOW_PB = _datetime_to_pb_timestamp(NOW)

        SERVE_NODES = 81

        client = _Client(self.PROJECT)
        instance = _Instance(self.INSTANCE_ID, client)
        cluster = self._make_one(self.CLUSTER_ID, instance,
                                 serve_nodes=SERVE_NODES)

        # Create request_pb
        request_pb = _ClusterPB(
            name=self.CLUSTER_NAME,
            serve_nodes=SERVE_NODES,
        )

        # Create response_pb
        OP_ID = 5678
        OP_NAME = (
            'operations/projects/%s/instances/%s/clusters/%s/operations/%d' %
            (self.PROJECT, self.INSTANCE_ID, self.CLUSTER_ID, OP_ID))
        metadata = messages_v2_pb2.UpdateClusterMetadata(request_time=NOW_PB)
        type_url = 'type.googleapis.com/%s' % (
            messages_v2_pb2.UpdateClusterMetadata.DESCRIPTOR.full_name,)
        response_pb = operations_pb2.Operation(
            name=OP_NAME,
            metadata=Any(
                type_url=type_url,
                value=metadata.SerializeToString()
            )
        )

        # Patch the stub used by the API method.
        client._instance_stub = stub = _FakeStub(response_pb)

        result = cluster.update()

        self.assertIsInstance(result, operation.Operation)
        self.assertEqual(result.operation.name, OP_NAME)
        self.assertIsInstance(result.metadata,
                              messages_v2_pb2.UpdateClusterMetadata)
        self.assertEqual(result.metadata.request_time, NOW_PB)

        self.assertEqual(len(stub.method_calls), 1)
        api_name, args, kwargs = stub.method_calls[0]
        self.assertEqual(api_name, 'UpdateCluster')
        request_pb, = args
        self.assertIsInstance(request_pb, data_v2_pb2.Cluster)
        self.assertEqual(request_pb.name, self.CLUSTER_NAME)
        self.assertEqual(request_pb.serve_nodes, SERVE_NODES)
        self.assertEqual(kwargs, {})
Esempio n. 58
0
    def test_commit(self):
        from tests.unit._testing import _FakeStub
        from google.cloud.bigtable.row_filters import RowSampleFilter

        row_key = b'row_key'
        table_name = 'projects/more-stuff'
        column_family_id1 = u'column_family_id1'
        column_family_id2 = u'column_family_id2'
        column_family_id3 = u'column_family_id3'
        column1 = b'column1'
        column2 = b'column2'
        client = _Client()
        table = _Table(table_name, client=client)
        row_filter = RowSampleFilter(0.33)
        row = self._make_one(row_key, table, filter_=row_filter)

        # Create request_pb
        value1 = b'bytes-value'
        mutation1 = _MutationPB(
            set_cell=_MutationSetCellPB(
                family_name=column_family_id1,
                column_qualifier=column1,
                timestamp_micros=-1,  # Default value.
                value=value1,
            ),
        )
        mutation2 = _MutationPB(
            delete_from_row=_MutationDeleteFromRowPB(),
        )
        mutation3 = _MutationPB(
            delete_from_column=_MutationDeleteFromColumnPB(
                family_name=column_family_id2,
                column_qualifier=column2,
            ),
        )
        mutation4 = _MutationPB(
            delete_from_family=_MutationDeleteFromFamilyPB(
                family_name=column_family_id3,
            ),
        )
        request_pb = _CheckAndMutateRowRequestPB(
            table_name=table_name,
            row_key=row_key,
            predicate_filter=row_filter.to_pb(),
            true_mutations=[mutation1, mutation3, mutation4],
            false_mutations=[mutation2],
        )

        # Create response_pb
        predicate_matched = True
        response_pb = _CheckAndMutateRowResponsePB(
            predicate_matched=predicate_matched)

        # Patch the stub used by the API method.
        client._data_stub = stub = _FakeStub(response_pb)

        # Create expected_result.
        expected_result = predicate_matched

        # Perform the method and check the result.
        row.set_cell(column_family_id1, column1, value1, state=True)
        row.delete(state=False)
        row.delete_cell(column_family_id2, column2, state=True)
        row.delete_cells(column_family_id3, row.ALL_COLUMNS, state=True)
        result = row.commit()
        self.assertEqual(result, expected_result)
        self.assertEqual(stub.method_calls, [(
            'CheckAndMutateRow',
            (request_pb,),
            {},
        )])
        self.assertEqual(row._true_pb_mutations, [])
        self.assertEqual(row._false_pb_mutations, [])