# Copyright 2016 Google LLC
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.


import unittest
import mock

from google.api_core.exceptions import DeadlineExceeded
from ._testing import _make_credentials
from google.cloud.bigtable.row_set import RowRange
from google.cloud.bigtable_v2.proto import (
    data_pb2 as data_v2_pb2)


class MultiCallableStub(object):
    """Stub for the grpc.UnaryUnaryMultiCallable interface."""

    def __init__(self, method, channel_stub):
        self.method = method
        self.channel_stub = channel_stub

    def __call__(self, request, timeout=None, metadata=None, credentials=None):
        self.channel_stub.requests.append((self.method, request))

        return self.channel_stub.responses.pop()


class ChannelStub(object):
    """Stub for the grpc.Channel interface."""

    def __init__(self, responses=[]):
        self.responses = responses
        self.requests = []

    def unary_unary(self,
                    method,
                    request_serializer=None,
                    response_deserializer=None):
        return MultiCallableStub(method, self)

    def unary_stream(self,
                     method,
                     request_serializer=None,
                     response_deserializer=None):
        return MultiCallableStub(method, self)


class TestCell(unittest.TestCase):
    timestamp_micros = 18738724000  # Make sure millis granularity

    @staticmethod
    def _get_target_class():
        from google.cloud.bigtable.row_data import Cell

        return Cell

    def _make_one(self, *args, **kwargs):
        return self._get_target_class()(*args, **kwargs)

    def _from_pb_test_helper(self, labels=None):
        import datetime
        from google.cloud._helpers import _EPOCH
        from google.cloud.bigtable_v2.proto import (
            data_pb2 as data_v2_pb2)

        timestamp_micros = TestCell.timestamp_micros
        timestamp = _EPOCH + datetime.timedelta(microseconds=timestamp_micros)
        value = b'value-bytes'

        if labels is None:
            cell_pb = data_v2_pb2.Cell(
                value=value, timestamp_micros=timestamp_micros)
            cell_expected = self._make_one(value, timestamp_micros)
        else:
            cell_pb = data_v2_pb2.Cell(
                value=value, timestamp_micros=timestamp_micros, labels=labels)
            cell_expected = self._make_one(
                value, timestamp_micros, labels=labels)

        klass = self._get_target_class()
        result = klass.from_pb(cell_pb)
        self.assertEqual(result, cell_expected)
        self.assertEqual(result.timestamp, timestamp)

    def test_from_pb(self):
        self._from_pb_test_helper()

    def test_from_pb_with_labels(self):
        labels = [u'label1', u'label2']
        self._from_pb_test_helper(labels)

    def test_constructor(self):
        value = object()
        cell = self._make_one(value, TestCell.timestamp_micros)
        self.assertEqual(cell.value, value)

    def test___eq__(self):
        value = object()
        cell1 = self._make_one(value, TestCell.timestamp_micros)
        cell2 = self._make_one(value, TestCell.timestamp_micros)
        self.assertEqual(cell1, cell2)

    def test___eq__type_differ(self):
        cell1 = self._make_one(None, None)
        cell2 = object()
        self.assertNotEqual(cell1, cell2)

    def test___ne__same_value(self):
        value = object()
        cell1 = self._make_one(value, TestCell.timestamp_micros)
        cell2 = self._make_one(value, TestCell.timestamp_micros)
        comparison_val = (cell1 != cell2)
        self.assertFalse(comparison_val)

    def test___ne__(self):
        value1 = 'value1'
        value2 = 'value2'
        cell1 = self._make_one(value1, TestCell.timestamp_micros)
        cell2 = self._make_one(value2, TestCell.timestamp_micros)
        self.assertNotEqual(cell1, cell2)


class TestPartialRowData(unittest.TestCase):

    @staticmethod
    def _get_target_class():
        from google.cloud.bigtable.row_data import PartialRowData

        return PartialRowData

    def _make_one(self, *args, **kwargs):
        return self._get_target_class()(*args, **kwargs)

    def test_constructor(self):
        row_key = object()
        partial_row_data = self._make_one(row_key)
        self.assertIs(partial_row_data._row_key, row_key)
        self.assertEqual(partial_row_data._cells, {})

    def test___eq__(self):
        row_key = object()
        partial_row_data1 = self._make_one(row_key)
        partial_row_data2 = self._make_one(row_key)
        self.assertEqual(partial_row_data1, partial_row_data2)

    def test___eq__type_differ(self):
        partial_row_data1 = self._make_one(None)
        partial_row_data2 = object()
        self.assertNotEqual(partial_row_data1, partial_row_data2)

    def test___ne__same_value(self):
        row_key = object()
        partial_row_data1 = self._make_one(row_key)
        partial_row_data2 = self._make_one(row_key)
        comparison_val = (partial_row_data1 != partial_row_data2)
        self.assertFalse(comparison_val)

    def test___ne__(self):
        row_key1 = object()
        partial_row_data1 = self._make_one(row_key1)
        row_key2 = object()
        partial_row_data2 = self._make_one(row_key2)
        self.assertNotEqual(partial_row_data1, partial_row_data2)

    def test___ne__cells(self):
        row_key = object()
        partial_row_data1 = self._make_one(row_key)
        partial_row_data1._cells = object()
        partial_row_data2 = self._make_one(row_key)
        self.assertNotEqual(partial_row_data1, partial_row_data2)

    def test_to_dict(self):
        cell1 = object()
        cell2 = object()
        cell3 = object()

        family_name1 = u'name1'
        family_name2 = u'name2'
        qual1 = b'col1'
        qual2 = b'col2'
        qual3 = b'col3'

        partial_row_data = self._make_one(None)
        partial_row_data._cells = {
            family_name1: {
                qual1: cell1,
                qual2: cell2,
            },
            family_name2: {
                qual3: cell3,
            },
        }

        result = partial_row_data.to_dict()
        expected_result = {
            b'name1:col1': cell1,
            b'name1:col2': cell2,
            b'name2:col3': cell3,
        }
        self.assertEqual(result, expected_result)

    def test_cell_value(self):
        family_name = u'name1'
        qualifier = b'col1'
        cell = _make_cell(b'value-bytes')

        partial_row_data = self._make_one(None)
        partial_row_data._cells = {
            family_name: {
                qualifier: [cell],
            },
        }

        result = partial_row_data.cell_value(family_name, qualifier)
        self.assertEqual(result, cell.value)

    def test_cell_value_invalid_index(self):
        family_name = u'name1'
        qualifier = b'col1'
        cell = _make_cell(b'')

        partial_row_data = self._make_one(None)
        partial_row_data._cells = {
            family_name: {
                qualifier: [cell],
            },
        }

        with self.assertRaises(IndexError):
            partial_row_data.cell_value(family_name, qualifier, index=None)

    def test_cell_value_invalid_column_family_key(self):
        family_name = u'name1'
        qualifier = b'col1'

        partial_row_data = self._make_one(None)

        with self.assertRaises(KeyError):
            partial_row_data.cell_value(family_name, qualifier)

    def test_cell_value_invalid_column_key(self):
        family_name = u'name1'
        qualifier = b'col1'

        partial_row_data = self._make_one(None)
        partial_row_data._cells = {
            family_name: {},
        }

        with self.assertRaises(KeyError):
            partial_row_data.cell_value(family_name, qualifier)

    def test_cell_values(self):
        family_name = u'name1'
        qualifier = b'col1'
        cell = _make_cell(b'value-bytes')

        partial_row_data = self._make_one(None)
        partial_row_data._cells = {
            family_name: {
                qualifier: [cell],
            },
        }

        values = []
        for value, timestamp_micros in partial_row_data.cell_values(
                family_name, qualifier):
            values.append(value)

        self.assertEqual(values[0], cell.value)

    def test_cell_values_with_max_count(self):
        family_name = u'name1'
        qualifier = b'col1'
        cell_1 = _make_cell(b'value-bytes-1')
        cell_2 = _make_cell(b'value-bytes-2')

        partial_row_data = self._make_one(None)
        partial_row_data._cells = {
            family_name: {
                qualifier: [cell_1, cell_2],
            },
        }

        values = []
        for value, timestamp_micros in partial_row_data.cell_values(
                family_name, qualifier, max_count=1):
            values.append(value)

        self.assertEqual(1, len(values))
        self.assertEqual(values[0], cell_1.value)

    def test_cells_property(self):
        partial_row_data = self._make_one(None)
        cells = {1: 2}
        partial_row_data._cells = cells
        # Make sure we get a copy, not the original.
        self.assertIsNot(partial_row_data.cells, cells)
        self.assertEqual(partial_row_data.cells, cells)

    def test_row_key_getter(self):
        row_key = object()
        partial_row_data = self._make_one(row_key)
        self.assertIs(partial_row_data.row_key, row_key)


class _Client(object):

    data_stub = None


class Test_retry_read_rows_exception(unittest.TestCase):

    @staticmethod
    def _call_fut(exc):
        from google.cloud.bigtable.row_data import _retry_read_rows_exception

        return _retry_read_rows_exception(exc)

    @staticmethod
    def _make_grpc_call_error(exception):
        from grpc import Call
        from grpc import RpcError

        class TestingException(Call, RpcError):
            def __init__(self, exception):
                self.exception = exception

            def code(self):
                return self.exception.grpc_status_code

            def details(self):
                return 'Testing'

        return TestingException(exception)

    def test_w_miss(self):
        from google.api_core.exceptions import Conflict

        exception = Conflict('testing')
        self.assertFalse(self._call_fut(exception))

    def test_w_service_unavailable(self):
        from google.api_core.exceptions import ServiceUnavailable

        exception = ServiceUnavailable('testing')
        self.assertTrue(self._call_fut(exception))

    def test_w_deadline_exceeded(self):
        from google.api_core.exceptions import DeadlineExceeded

        exception = DeadlineExceeded('testing')
        self.assertTrue(self._call_fut(exception))

    def test_w_miss_wrapped_in_grpc(self):
        from google.api_core.exceptions import Conflict

        wrapped = Conflict('testing')
        exception = self._make_grpc_call_error(wrapped)
        self.assertFalse(self._call_fut(exception))

    def test_w_service_unavailable_wrapped_in_grpc(self):
        from google.api_core.exceptions import ServiceUnavailable

        wrapped = ServiceUnavailable('testing')
        exception = self._make_grpc_call_error(wrapped)
        self.assertTrue(self._call_fut(exception))

    def test_w_deadline_exceeded_wrapped_in_grpc(self):
        from google.api_core.exceptions import DeadlineExceeded

        wrapped = DeadlineExceeded('testing')
        exception = self._make_grpc_call_error(wrapped)
        self.assertTrue(self._call_fut(exception))


class TestPartialRowsData(unittest.TestCase):
    ROW_KEY = b'row-key'
    FAMILY_NAME = u'family'
    QUALIFIER = b'qualifier'
    TIMESTAMP_MICROS = 100
    VALUE = b'value'

    @staticmethod
    def _get_target_class():
        from google.cloud.bigtable.row_data import PartialRowsData

        return PartialRowsData

    @staticmethod
    def _get_target_client_class():
        from google.cloud.bigtable.client import Client

        return Client

    def _make_one(self, *args, **kwargs):
        return self._get_target_class()(*args, **kwargs)

    def test_constructor(self):
        from google.cloud.bigtable.row_data import DEFAULT_RETRY_READ_ROWS
        client = _Client()
        client._data_stub = mock.MagicMock()
        request = object()
        partial_rows_data = self._make_one(client._data_stub.ReadRows,
                                           request)
        self.assertIs(partial_rows_data.request, request)
        self.assertEqual(partial_rows_data.rows, {})
        self.assertEqual(partial_rows_data.retry,
                         DEFAULT_RETRY_READ_ROWS)

    def test_constructor_with_retry(self):
        client = _Client()
        client._data_stub = mock.MagicMock()
        request = retry = object()
        partial_rows_data = self._make_one(client._data_stub.ReadRows,
                                           request, retry)
        self.assertIs(partial_rows_data.request, request)
        self.assertEqual(partial_rows_data.rows, {})
        self.assertEqual(partial_rows_data.retry,
                         retry)

    def test___eq__(self):
        client = _Client()
        client._data_stub = mock.MagicMock()
        request = object()
        partial_rows_data1 = self._make_one(client._data_stub.ReadRows,
                                            request)
        partial_rows_data2 = self._make_one(client._data_stub.ReadRows,
                                            request)
        self.assertEqual(partial_rows_data1.rows, partial_rows_data2.rows)

    def test___eq__type_differ(self):
        client = _Client()
        client._data_stub = mock.MagicMock()
        request = object()
        partial_rows_data1 = self._make_one(client._data_stub.ReadRows,
                                            request)
        partial_rows_data2 = object()
        self.assertNotEqual(partial_rows_data1, partial_rows_data2)

    def test___ne__same_value(self):
        client = _Client()
        client._data_stub = mock.MagicMock()
        request = object()
        partial_rows_data1 = self._make_one(client._data_stub.ReadRows,
                                            request)
        partial_rows_data2 = self._make_one(client._data_stub.ReadRows,
                                            request)
        comparison_val = (partial_rows_data1 != partial_rows_data2)
        self.assertTrue(comparison_val)

    def test___ne__(self):
        client = _Client()
        client._data_stub = mock.MagicMock()
        request = object()
        partial_rows_data1 = self._make_one(client._data_stub.ReadRows,
                                            request)
        partial_rows_data2 = self._make_one(client._data_stub.ReadRows,
                                            request)
        self.assertNotEqual(partial_rows_data1, partial_rows_data2)

    def test_rows_getter(self):
        client = _Client()
        client._data_stub = mock.MagicMock()
        request = object()
        partial_rows_data = self._make_one(client._data_stub.ReadRows,
                                           request)
        partial_rows_data.rows = value = object()
        self.assertIs(partial_rows_data.rows, value)

    def _make_client(self, *args, **kwargs):
        return self._get_target_client_class()(*args, **kwargs)

    def test_state_start(self):
        client = _Client()
        iterator = _MockCancellableIterator()
        client._data_stub = mock.MagicMock()
        client._data_stub.ReadRows.side_effect = [iterator]
        request = object()
        yrd = self._make_one(client._data_stub.ReadRows, request)
        self.assertEqual(yrd.state, yrd.NEW_ROW)

    def test_state_new_row_w_row(self):
        from google.cloud.bigtable_v2.gapic import bigtable_client

        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)
        iterator = _MockCancellableIterator(response)
        channel = ChannelStub(responses=[iterator])
        data_api = bigtable_client.BigtableClient(channel=channel)
        credentials = _make_credentials()
        client = self._make_client(project='project-id',
                                   credentials=credentials, admin=True)
        client._table_data_client = data_api
        request = object()

        yrd = self._make_one(
            client._table_data_client.transport.read_rows, request)

        yrd._response_iterator = iterator
        rows = [row for row in yrd]

        result = rows[0]
        self.assertEqual(result.row_key, self.ROW_KEY)
        self.assertEqual(yrd._counter, 1)
        self.assertEqual(yrd.state, yrd.NEW_ROW)

    def test_multiple_chunks(self):
        from google.cloud.bigtable_v2.gapic import bigtable_client

        chunk1 = _ReadRowsResponseCellChunkPB(
            row_key=self.ROW_KEY,
            family_name=self.FAMILY_NAME,
            qualifier=self.QUALIFIER,
            timestamp_micros=self.TIMESTAMP_MICROS,
            value=self.VALUE,
            commit_row=False,
        )
        chunk2 = _ReadRowsResponseCellChunkPB(
            qualifier=self.QUALIFIER + b'1',
            timestamp_micros=self.TIMESTAMP_MICROS,
            value=self.VALUE,
            commit_row=True,
        )
        chunks = [chunk1, chunk2]

        response = _ReadRowsResponseV2(chunks)
        iterator = _MockCancellableIterator(response)
        channel = ChannelStub(responses=[iterator])
        data_api = bigtable_client.BigtableClient(channel=channel)
        credentials = _make_credentials()
        client = self._make_client(project='project-id',
                                   credentials=credentials, admin=True)
        client._table_data_client = data_api
        request = object()

        yrd = self._make_one(
            client._table_data_client.transport.read_rows, request)

        yrd._response_iterator = iterator
        rows = [row for row in yrd]
        result = rows[0]
        self.assertEqual(result.row_key, self.ROW_KEY)
        self.assertEqual(yrd._counter, 1)
        self.assertEqual(yrd.state, yrd.NEW_ROW)

    def test_cancel(self):
        client = _Client()
        response_iterator = _MockCancellableIterator()
        client._data_stub = mock.MagicMock()
        client._data_stub.ReadRows.side_effect = [response_iterator]
        request = object()
        yield_rows_data = self._make_one(client._data_stub.ReadRows, request)
        self.assertEqual(response_iterator.cancel_calls, 0)
        yield_rows_data.cancel()
        self.assertEqual(response_iterator.cancel_calls, 1)

    # 'consume_next' tested via 'TestPartialRowsData_JSON_acceptance_tests'

    def test__copy_from_previous_unset(self):
        client = _Client()
        client._data_stub = mock.MagicMock()
        request = object()
        yrd = self._make_one(client._data_stub.ReadRows, request)
        cell = _PartialCellData()
        yrd._copy_from_previous(cell)
        self.assertEqual(cell.row_key, b'')
        self.assertEqual(cell.family_name, u'')
        self.assertIsNone(cell.qualifier)
        self.assertEqual(cell.timestamp_micros, 0)
        self.assertEqual(cell.labels, [])

    def test__copy_from_previous_blank(self):
        ROW_KEY = 'RK'
        FAMILY_NAME = u'A'
        QUALIFIER = b'C'
        TIMESTAMP_MICROS = 100
        LABELS = ['L1', 'L2']
        client = _Client()
        client._data_stub = mock.MagicMock()
        request = object()
        yrd = self._make_one(client._data_stub.ReadRows, request)
        cell = _PartialCellData(
            row_key=ROW_KEY,
            family_name=FAMILY_NAME,
            qualifier=QUALIFIER,
            timestamp_micros=TIMESTAMP_MICROS,
            labels=LABELS,
        )
        yrd._previous_cell = _PartialCellData()
        yrd._copy_from_previous(cell)
        self.assertEqual(cell.row_key, ROW_KEY)
        self.assertEqual(cell.family_name, FAMILY_NAME)
        self.assertEqual(cell.qualifier, QUALIFIER)
        self.assertEqual(cell.timestamp_micros, TIMESTAMP_MICROS)
        self.assertEqual(cell.labels, LABELS)

    def test__copy_from_previous_filled(self):
        ROW_KEY = 'RK'
        FAMILY_NAME = u'A'
        QUALIFIER = b'C'
        TIMESTAMP_MICROS = 100
        LABELS = ['L1', 'L2']
        client = _Client()
        client._data_stub = mock.MagicMock()
        request = object()
        yrd = self._make_one(client._data_stub.ReadRows, request)
        yrd._previous_cell = _PartialCellData(
            row_key=ROW_KEY,
            family_name=FAMILY_NAME,
            qualifier=QUALIFIER,
            timestamp_micros=TIMESTAMP_MICROS,
            labels=LABELS,
        )
        cell = _PartialCellData()
        yrd._copy_from_previous(cell)
        self.assertEqual(cell.row_key, ROW_KEY)
        self.assertEqual(cell.family_name, FAMILY_NAME)
        self.assertEqual(cell.qualifier, QUALIFIER)
        self.assertEqual(cell.timestamp_micros, 0)
        self.assertEqual(cell.labels, [])

    def test_valid_last_scanned_row_key_on_start(self):
        client = _Client()
        response = _ReadRowsResponseV2(
            chunks=(), last_scanned_row_key='2.AFTER')
        iterator = _MockCancellableIterator(response)
        client._data_stub = mock.MagicMock()
        client._data_stub.ReadRows.side_effect = [iterator]
        request = object()
        yrd = self._make_one(client._data_stub.ReadRows, request)
        yrd.last_scanned_row_key = '1.BEFORE'
        self._consume_all(yrd)
        self.assertEqual(yrd.last_scanned_row_key, '2.AFTER')

    def test_invalid_empty_chunk(self):
        from google.cloud.bigtable.row_data import InvalidChunk

        client = _Client()
        chunks = _generate_cell_chunks([''])
        response = _ReadRowsResponseV2(chunks)
        iterator = _MockCancellableIterator(response)
        client._data_stub = mock.MagicMock()
        client._data_stub.ReadRows.side_effect = [iterator]
        request = object()
        yrd = self._make_one(client._data_stub.ReadRows, request)
        with self.assertRaises(InvalidChunk):
            self._consume_all(yrd)

    def test_state_cell_in_progress(self):
        LABELS = ['L1', 'L2']

        request = object()
        read_rows = mock.MagicMock()
        yrd = self._make_one(read_rows, request)

        chunk = _ReadRowsResponseCellChunkPB(
            row_key=self.ROW_KEY,
            family_name=self.FAMILY_NAME,
            qualifier=self.QUALIFIER,
            timestamp_micros=self.TIMESTAMP_MICROS,
            value=self.VALUE,
            labels=LABELS
        )
        yrd._update_cell(chunk)

        more_cell_data = _ReadRowsResponseCellChunkPB(value=self.VALUE)
        yrd._update_cell(more_cell_data)

        self.assertEqual(yrd._cell.row_key, self.ROW_KEY)
        self.assertEqual(yrd._cell.family_name, self.FAMILY_NAME)
        self.assertEqual(yrd._cell.qualifier, self.QUALIFIER)
        self.assertEqual(yrd._cell.timestamp_micros,
                         self.TIMESTAMP_MICROS)
        self.assertEqual(yrd._cell.labels, LABELS)
        self.assertEqual(yrd._cell.value, self.VALUE + self.VALUE)

    def test_yield_rows_data(self):
        client = _Client()

        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)
        iterator = _MockCancellableIterator(response)
        client._data_stub = mock.MagicMock()
        client._data_stub.ReadRows.side_effect = [iterator]

        request = object()

        yrd = self._make_one(client._data_stub.ReadRows, request)

        result = self._consume_all(yrd)[0]

        self.assertEqual(result, self.ROW_KEY)

    def test_yield_retry_rows_data(self):
        from google.api_core import retry
        client = _Client()

        retry_read_rows = retry.Retry(
            predicate=_read_rows_retry_exception,
        )

        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)
        failure_iterator = _MockFailureIterator_1()
        iterator = _MockCancellableIterator(response)
        client._data_stub = mock.MagicMock()
        client._data_stub.ReadRows.side_effect = [failure_iterator,
                                                  iterator]

        request = object()

        yrd = self._make_one(client._data_stub.ReadRows, request,
                             retry_read_rows)

        result = self._consume_all(yrd)[0]

        self.assertEqual(result, self.ROW_KEY)

    def _consume_all(self, yrd):
        return [row.row_key for row in yrd]


class Test_ReadRowsRequestManager(unittest.TestCase):

    @classmethod
    def setUpClass(cls):
        cls.table_name = 'table_name'
        cls.row_range1 = RowRange(b"row_key21", b"row_key29")
        cls.row_range2 = RowRange(b"row_key31", b"row_key39")
        cls.row_range3 = RowRange(b"row_key41", b"row_key49")

        cls.request = _ReadRowsRequestPB(table_name=cls.table_name)
        cls.request.rows.row_ranges.add(**cls.row_range1.get_range_kwargs())
        cls.request.rows.row_ranges.add(**cls.row_range2.get_range_kwargs())
        cls.request.rows.row_ranges.add(**cls.row_range3.get_range_kwargs())

    @staticmethod
    def _get_target_class():
        from google.cloud.bigtable.row_data import _ReadRowsRequestManager
        return _ReadRowsRequestManager

    def _make_one(self, *args, **kwargs):
        return self._get_target_class()(*args, **kwargs)

    def test_constructor(self):
        request = mock.Mock()
        last_scanned_key = "last_key"
        rows_read_so_far = 10

        request_manager = self._make_one(request, last_scanned_key,
                                         rows_read_so_far)
        self.assertEqual(request, request_manager.message)
        self.assertEqual(last_scanned_key, request_manager.last_scanned_key)
        self.assertEqual(rows_read_so_far, request_manager.rows_read_so_far)

    def test__filter_row_key(self):
        table_name = 'table_name'
        request = _ReadRowsRequestPB(table_name=table_name)
        request.rows.row_keys.extend([b'row_key1', b'row_key2',
                                      b'row_key3', b'row_key4'])

        last_scanned_key = b"row_key2"
        request_manager = self._make_one(request, last_scanned_key, 2)
        row_keys = request_manager._filter_rows_keys()

        expected_row_keys = [b'row_key3', b'row_key4']
        self.assertEqual(expected_row_keys, row_keys)

    def test__filter_row_ranges_all_ranges_added_back(self):
        last_scanned_key = b"row_key14"
        request_manager = self._make_one(self.request, last_scanned_key, 2)
        row_ranges = request_manager._filter_row_ranges()

        exp_row_range1 = data_v2_pb2.RowRange(start_key_closed=b"row_key21",
                                              end_key_open=b"row_key29")
        exp_row_range2 = data_v2_pb2.RowRange(start_key_closed=b"row_key31",
                                              end_key_open=b"row_key39")
        exp_row_range3 = data_v2_pb2.RowRange(start_key_closed=b"row_key41",
                                              end_key_open=b"row_key49")
        exp_row_ranges = [exp_row_range1, exp_row_range2, exp_row_range3]

        self.assertEqual(exp_row_ranges, row_ranges)

    def test__filter_row_ranges_all_ranges_already_read(self):
        last_scanned_key = b"row_key54"
        request_manager = self._make_one(self.request, last_scanned_key, 2)
        row_ranges = request_manager._filter_row_ranges()

        self.assertEqual(row_ranges, [])

    def test__filter_row_ranges_all_ranges_already_read_open_closed(self):
        last_scanned_key = b"row_key54"

        row_range1 = RowRange(b"row_key21", b"row_key29", False, True)
        row_range2 = RowRange(b"row_key31", b"row_key39")
        row_range3 = RowRange(b"row_key41", b"row_key49", False, True)

        request = _ReadRowsRequestPB(table_name=self.table_name)
        request.rows.row_ranges.add(**row_range1.get_range_kwargs())
        request.rows.row_ranges.add(**row_range2.get_range_kwargs())
        request.rows.row_ranges.add(**row_range3.get_range_kwargs())

        request_manager = self._make_one(request, last_scanned_key, 2)
        request_manager.new_message = _ReadRowsRequestPB(
            table_name=self.table_name)
        row_ranges = request_manager._filter_row_ranges()

        self.assertEqual(row_ranges, [])

    def test__filter_row_ranges_some_ranges_already_read(self):
        last_scanned_key = b"row_key22"
        request_manager = self._make_one(self.request, last_scanned_key, 2)
        request_manager.new_message = _ReadRowsRequestPB(
            table_name=self.table_name)
        row_ranges = request_manager._filter_row_ranges()

        exp_row_range1 = data_v2_pb2.RowRange(start_key_open=b"row_key22",
                                              end_key_open=b"row_key29")
        exp_row_range2 = data_v2_pb2.RowRange(start_key_closed=b"row_key31",
                                              end_key_open=b"row_key39")
        exp_row_range3 = data_v2_pb2.RowRange(start_key_closed=b"row_key41",
                                              end_key_open=b"row_key49")
        exp_row_ranges = [exp_row_range1, exp_row_range2, exp_row_range3]

        self.assertEqual(exp_row_ranges, row_ranges)

    def test_build_updated_request(self):
        from google.cloud.bigtable.row_filters import RowSampleFilter
        row_filter = RowSampleFilter(0.33)
        last_scanned_key = b"row_key25"
        request = _ReadRowsRequestPB(filter=row_filter.to_pb(),
                                     rows_limit=8,
                                     table_name=self.table_name)
        request.rows.row_ranges.add(**self.row_range1.get_range_kwargs())

        request_manager = self._make_one(request, last_scanned_key, 2)

        result = request_manager.build_updated_request()

        expected_result = _ReadRowsRequestPB(table_name=self.table_name,
                                             filter=row_filter.to_pb(),
                                             rows_limit=6)
        expected_result.rows.row_ranges.add(
            start_key_open=last_scanned_key,
            end_key_open=self.row_range1.end_key
        )

        self.assertEqual(expected_result, result)

    def test_build_updated_request_full_table(self):
        last_scanned_key = b"row_key14"

        request = _ReadRowsRequestPB(table_name=self.table_name)
        request_manager = self._make_one(request, last_scanned_key, 2)

        result = request_manager.build_updated_request()
        expected_result = _ReadRowsRequestPB(table_name=self.table_name,
                                             filter={})
        expected_result.rows.row_ranges.add(start_key_open=last_scanned_key)
        self.assertEqual(expected_result, result)

    def test_build_updated_request_no_start_key(self):
        from google.cloud.bigtable.row_filters import RowSampleFilter
        row_filter = RowSampleFilter(0.33)
        last_scanned_key = b"row_key25"
        request = _ReadRowsRequestPB(filter=row_filter.to_pb(),
                                     rows_limit=8,
                                     table_name=self.table_name)
        request.rows.row_ranges.add(end_key_open=b"row_key29")

        request_manager = self._make_one(request, last_scanned_key, 2)

        result = request_manager.build_updated_request()

        expected_result = _ReadRowsRequestPB(table_name=self.table_name,
                                             filter=row_filter.to_pb(),
                                             rows_limit=6)
        expected_result.rows.row_ranges.add(start_key_open=last_scanned_key,
                                            end_key_open=b"row_key29")

        self.assertEqual(expected_result, result)

    def test_build_updated_request_no_end_key(self):
        from google.cloud.bigtable.row_filters import RowSampleFilter
        row_filter = RowSampleFilter(0.33)
        last_scanned_key = b"row_key25"
        request = _ReadRowsRequestPB(filter=row_filter.to_pb(),
                                     rows_limit=8,
                                     table_name=self.table_name)
        request.rows.row_ranges.add(start_key_closed=b"row_key20")

        request_manager = self._make_one(request, last_scanned_key, 2)

        result = request_manager.build_updated_request()

        expected_result = _ReadRowsRequestPB(table_name=self.table_name,
                                             filter=row_filter.to_pb(),
                                             rows_limit=6)
        expected_result.rows.row_ranges.add(start_key_open=last_scanned_key)

        self.assertEqual(expected_result, result)

    def test_build_updated_request_rows(self):
        from google.cloud.bigtable.row_filters import RowSampleFilter
        row_filter = RowSampleFilter(0.33)
        last_scanned_key = b"row_key4"
        request = _ReadRowsRequestPB(filter=row_filter.to_pb(),
                                     rows_limit=5,
                                     table_name=self.table_name)
        request.rows.row_keys.extend([b'row_key1', b'row_key2',
                                      b'row_key4', b'row_key5',
                                      b'row_key7', b'row_key9'])

        request_manager = self._make_one(request, last_scanned_key, 3)

        result = request_manager.build_updated_request()

        expected_result = _ReadRowsRequestPB(table_name=self.table_name,
                                             filter=row_filter.to_pb(),
                                             rows_limit=2)
        expected_result.rows.row_keys.extend([b'row_key5', b'row_key7',
                                              b'row_key9'])

        self.assertEqual(expected_result, result)

    def test_build_updated_request_rows_limit(self):
        last_scanned_key = b"row_key14"

        request = _ReadRowsRequestPB(table_name=self.table_name, rows_limit=10)
        request_manager = self._make_one(request, last_scanned_key, 2)

        result = request_manager.build_updated_request()
        expected_result = _ReadRowsRequestPB(table_name=self.table_name,
                                             filter={},
                                             rows_limit=8)
        expected_result.rows.row_ranges.add(start_key_open=last_scanned_key)
        self.assertEqual(expected_result, result)

    def test__key_already_read(self):
        last_scanned_key = b"row_key14"
        request = _ReadRowsRequestPB(table_name=self.table_name)
        request_manager = self._make_one(request, last_scanned_key, 2)

        self.assertTrue(request_manager._key_already_read(b"row_key11"))
        self.assertFalse(request_manager._key_already_read(b"row_key16"))


class TestPartialRowsData_JSON_acceptance_tests(unittest.TestCase):

    _json_tests = None

    @staticmethod
    def _get_target_class():
        from google.cloud.bigtable.row_data import PartialRowsData

        return PartialRowsData

    def _make_one(self, *args, **kwargs):
        return self._get_target_class()(*args, **kwargs)

    def _load_json_test(self, test_name):
        import os

        if self.__class__._json_tests is None:
            dirname = os.path.dirname(__file__)
            filename = os.path.join(dirname, 'read-rows-acceptance-test.json')
            raw = _parse_readrows_acceptance_tests(filename)
            tests = self.__class__._json_tests = {}
            for (name, chunks, results) in raw:
                tests[name] = chunks, results
        return self.__class__._json_tests[test_name]

    # JSON Error cases:  invalid chunks

    def _fail_during_consume(self, testcase_name):
        from google.cloud.bigtable.row_data import InvalidChunk

        client = _Client()
        chunks, results = self._load_json_test(testcase_name)
        response = _ReadRowsResponseV2(chunks)
        iterator = _MockCancellableIterator(response)
        client._data_stub = mock.MagicMock()
        client._data_stub.ReadRows.side_effect = [iterator]
        request = object()
        prd = self._make_one(client._data_stub.ReadRows, request)
        with self.assertRaises(InvalidChunk):
            prd.consume_all()
        expected_result = self._sort_flattend_cells(
            [result for result in results if not result['error']])
        flattened = self._sort_flattend_cells(_flatten_cells(prd))
        self.assertEqual(flattened, expected_result)

    def test_invalid_no_cell_key_before_commit(self):
        self._fail_during_consume('invalid - no cell key before commit')

    def test_invalid_no_cell_key_before_value(self):
        self._fail_during_consume('invalid - no cell key before value')

    def test_invalid_new_col_family_wo_qualifier(self):
        self._fail_during_consume(
            'invalid - new col family must specify qualifier')

    def test_invalid_no_commit_between_rows(self):
        self._fail_during_consume('invalid - no commit between rows')

    def test_invalid_no_commit_after_first_row(self):
        self._fail_during_consume('invalid - no commit after first row')

    def test_invalid_duplicate_row_key(self):
        self._fail_during_consume('invalid - duplicate row key')

    def test_invalid_new_row_missing_row_key(self):
        self._fail_during_consume('invalid - new row missing row key')

    def test_invalid_bare_reset(self):
        self._fail_during_consume('invalid - bare reset')

    def test_invalid_bad_reset_no_commit(self):
        self._fail_during_consume('invalid - bad reset, no commit')

    def test_invalid_missing_key_after_reset(self):
        self._fail_during_consume('invalid - missing key after reset')

    def test_invalid_reset_with_chunk(self):
        self._fail_during_consume('invalid - reset with chunk')

    def test_invalid_commit_with_chunk(self):
        self._fail_during_consume('invalid - commit with chunk')

    # JSON Error cases:  incomplete final row

    def _sort_flattend_cells(self, flattened):
        import operator

        key_func = operator.itemgetter('rk', 'fm', 'qual')
        return sorted(flattened, key=key_func)

    def _incomplete_final_row(self, testcase_name):
        client = _Client()
        chunks, results = self._load_json_test(testcase_name)
        response = _ReadRowsResponseV2(chunks)
        iterator = _MockCancellableIterator(response)
        client._data_stub = mock.MagicMock()
        client._data_stub.ReadRows.side_effect = [iterator]
        request = object()
        prd = self._make_one(client._data_stub.ReadRows, request)
        with self.assertRaises(ValueError):
            prd.consume_all()
        self.assertEqual(prd.state, prd.ROW_IN_PROGRESS)
        expected_result = self._sort_flattend_cells(
            [result for result in results if not result['error']])
        flattened = self._sort_flattend_cells(_flatten_cells(prd))
        self.assertEqual(flattened, expected_result)

    def test_invalid_no_commit(self):
        self._incomplete_final_row('invalid - no commit')

    def test_invalid_last_row_missing_commit(self):
        self._incomplete_final_row('invalid - last row missing commit')

    # Non-error cases

    _marker = object()

    def _match_results(self, testcase_name, expected_result=_marker):
        client = _Client()
        chunks, results = self._load_json_test(testcase_name)
        response = _ReadRowsResponseV2(chunks)
        iterator = _MockCancellableIterator(response)
        client._data_stub = mock.MagicMock()
        client._data_stub.ReadRows.side_effect = [iterator]
        request = object()
        prd = self._make_one(client._data_stub.ReadRows, request)
        prd.consume_all()
        flattened = self._sort_flattend_cells(_flatten_cells(prd))
        if expected_result is self._marker:
            expected_result = self._sort_flattend_cells(results)
        self.assertEqual(flattened, expected_result)

    def test_bare_commit_implies_ts_zero(self):
        self._match_results('bare commit implies ts=0')

    def test_simple_row_with_timestamp(self):
        self._match_results('simple row with timestamp')

    def test_missing_timestamp_implies_ts_zero(self):
        self._match_results('missing timestamp, implied ts=0')

    def test_empty_cell_value(self):
        self._match_results('empty cell value')

    def test_two_unsplit_cells(self):
        self._match_results('two unsplit cells')

    def test_two_qualifiers(self):
        self._match_results('two qualifiers')

    def test_two_families(self):
        self._match_results('two families')

    def test_with_labels(self):
        self._match_results('with labels')

    def test_split_cell_bare_commit(self):
        self._match_results('split cell, bare commit')

    def test_split_cell(self):
        self._match_results('split cell')

    def test_split_four_ways(self):
        self._match_results('split four ways')

    def test_two_split_cells(self):
        self._match_results('two split cells')

    def test_multi_qualifier_splits(self):
        self._match_results('multi-qualifier splits')

    def test_multi_qualifier_multi_split(self):
        self._match_results('multi-qualifier multi-split')

    def test_multi_family_split(self):
        self._match_results('multi-family split')

    def test_two_rows(self):
        self._match_results('two rows')

    def test_two_rows_implicit_timestamp(self):
        self._match_results('two rows implicit timestamp')

    def test_two_rows_empty_value(self):
        self._match_results('two rows empty value')

    def test_two_rows_one_with_multiple_cells(self):
        self._match_results('two rows, one with multiple cells')

    def test_two_rows_multiple_cells_multiple_families(self):
        self._match_results('two rows, multiple cells, multiple families')

    def test_two_rows_multiple_cells(self):
        self._match_results('two rows, multiple cells')

    def test_two_rows_four_cells_two_labels(self):
        self._match_results('two rows, four cells, 2 labels')

    def test_two_rows_with_splits_same_timestamp(self):
        self._match_results('two rows with splits, same timestamp')

    def test_no_data_after_reset(self):
        # JSON testcase has `"results": null`
        self._match_results('no data after reset', expected_result=[])

    def test_simple_reset(self):
        self._match_results('simple reset')

    def test_reset_to_new_val(self):
        self._match_results('reset to new val')

    def test_reset_to_new_qual(self):
        self._match_results('reset to new qual')

    def test_reset_with_splits(self):
        self._match_results('reset with splits')

    def test_two_resets(self):
        self._match_results('two resets')

    def test_reset_to_new_row(self):
        self._match_results('reset to new row')

    def test_reset_in_between_chunks(self):
        self._match_results('reset in between chunks')

    def test_empty_cell_chunk(self):
        self._match_results('empty cell chunk')

    def test_empty_second_qualifier(self):
        self._match_results('empty second qualifier')


def _flatten_cells(prd):
    # Match results format from JSON testcases.
    # Doesn't handle error cases.
    from google.cloud._helpers import _bytes_to_unicode
    from google.cloud._helpers import _microseconds_from_datetime

    for row_key, row in prd.rows.items():
        for family_name, family in row.cells.items():
            for qualifier, column in family.items():
                for cell in column:
                    yield {
                        u'rk': _bytes_to_unicode(row_key),
                        u'fm': family_name,
                        u'qual': _bytes_to_unicode(qualifier),
                        u'ts': _microseconds_from_datetime(cell.timestamp),
                        u'value': _bytes_to_unicode(cell.value),
                        u'label': u' '.join(cell.labels),
                        u'error': False,
                    }


class _MockCancellableIterator(object):

    cancel_calls = 0

    def __init__(self, *values):
        self.iter_values = iter(values)

    def cancel(self):
        self.cancel_calls += 1

    def next(self):
        return next(self.iter_values)

    __next__ = next


class _MockFailureIterator_1(object):

    def next(self):
        raise DeadlineExceeded("Failed to read from server")

    __next__ = next


class _PartialCellData(object):

    row_key = b''
    family_name = u''
    qualifier = None
    timestamp_micros = 0

    def __init__(self, **kw):
        self.labels = kw.pop('labels', [])
        self.__dict__.update(kw)


class _ReadRowsResponseV2(object):

    def __init__(self, chunks, last_scanned_row_key=''):
        self.chunks = chunks
        self.last_scanned_row_key = last_scanned_row_key


def _generate_cell_chunks(chunk_text_pbs):
    from google.protobuf.text_format import Merge
    from google.cloud.bigtable_v2.proto.bigtable_pb2 import ReadRowsResponse

    chunks = []

    for chunk_text_pb in chunk_text_pbs:
        chunk = ReadRowsResponse.CellChunk()
        chunks.append(Merge(chunk_text_pb, chunk))

    return chunks


def _parse_readrows_acceptance_tests(filename):
    """Parse acceptance tests from JSON

    See
    https://github.com/GoogleCloudPlatform/cloud-bigtable-client/blob/\
    4d3185662ca61bc9fa1bdf1ec0166f6e5ecf86c6/bigtable-client-core/src/\
    test/resources/com/google/cloud/bigtable/grpc/scanner/v2/
    read-rows-acceptance-test.json
    """
    import json

    with open(filename) as json_file:
        test_json = json.load(json_file)

    for test in test_json['tests']:
        name = test['name']
        chunks = _generate_cell_chunks(test['chunks'])
        results = test['results']
        yield name, chunks, results


def _ReadRowsResponseCellChunkPB(*args, **kw):
    from google.cloud.bigtable_v2.proto import (
        bigtable_pb2 as messages_v2_pb2)

    family_name = kw.pop('family_name', None)
    qualifier = kw.pop('qualifier', None)
    message = messages_v2_pb2.ReadRowsResponse.CellChunk(*args, **kw)

    if family_name:
        message.family_name.value = family_name
    if qualifier:
        message.qualifier.value = qualifier

    return message


def _make_cell(value):
    from google.cloud.bigtable import row_data

    return row_data.Cell(value, TestCell.timestamp_micros)


def _ReadRowsRequestPB(*args, **kw):
    from google.cloud.bigtable_v2.proto import (
        bigtable_pb2 as messages_v2_pb2)

    return messages_v2_pb2.ReadRowsRequest(*args, **kw)


def _read_rows_retry_exception(exc):
    return isinstance(exc, DeadlineExceeded)
