# Copyright 2018 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 ._testing import _make_credentials

from google.cloud.bigtable.batcher import MutationsBatcher

from google.cloud.bigtable.row import DirectRow

class TestMutationsBatcher(unittest.TestCase):

from grpc import StatusCode

TABLE_ID = "table-id"

TABLE_NAME = "/tables/" + TABLE_ID

# RPC Status Codes

SUCCESS = StatusCode.OK.value[0]

@staticmethod

def _get_target_class():

from google.cloud.bigtable.table import Table

return Table

def _make_table(self, *args, **kwargs):

return self._get_target_class()(*args, **kwargs)

@staticmethod

def _get_target_client_class():

from google.cloud.bigtable.client import Client

return Client

def _make_client(self, *args, **kwargs):

return self._get_target_client_class()(*args, **kwargs)

def test_constructor(self):

credentials = _make_credentials()

client = self._make_client(

project="project-id", credentials=credentials, admin=True

)

instance = client.instance(instance_id="instance-id")

table = self._make_table(self.TABLE_ID, instance)

mutation_batcher = MutationsBatcher(table)

self.assertEqual(table, mutation_batcher.table)

def test_mutate_row(self):

table = _Table(self.TABLE_NAME)

mutation_batcher = MutationsBatcher(table=table)

rows = [

DirectRow(row_key=b"row_key"),

DirectRow(row_key=b"row_key_2"),

DirectRow(row_key=b"row_key_3"),

DirectRow(row_key=b"row_key_4"),

]

mutation_batcher.mutate_rows(rows)

mutation_batcher.flush()

self.assertEqual(table.mutation_calls, 1)

def test_mutate_rows(self):

table = _Table(self.TABLE_NAME)

mutation_batcher = MutationsBatcher(table=table)

row = DirectRow(row_key=b"row_key")

row.set_cell("cf1", b"c1", 1)

row.set_cell("cf1", b"c2", 2)

row.set_cell("cf1", b"c3", 3)

row.set_cell("cf1", b"c4", 4)

mutation_batcher.mutate(row)

mutation_batcher.flush()

self.assertEqual(table.mutation_calls, 1)

def test_flush_with_no_rows(self):

table = _Table(self.TABLE_NAME)

mutation_batcher = MutationsBatcher(table=table)

mutation_batcher.flush()

self.assertEqual(table.mutation_calls, 0)

def test_add_row_with_max_flush_count(self):

table = _Table(self.TABLE_NAME)

mutation_batcher = MutationsBatcher(table=table, flush_count=3)

row_1 = DirectRow(row_key=b"row_key_1")

row_2 = DirectRow(row_key=b"row_key_2")

row_3 = DirectRow(row_key=b"row_key_3")

mutation_batcher.mutate(row_1)

mutation_batcher.mutate(row_2)

mutation_batcher.mutate(row_3)

self.assertEqual(table.mutation_calls, 1)

@mock.patch("google.cloud.bigtable.batcher.MAX_MUTATIONS", new=3)

def test_mutate_row_with_max_mutations_failure(self):

from google.cloud.bigtable.batcher import MaxMutationsError

table = _Table(self.TABLE_NAME)

mutation_batcher = MutationsBatcher(table=table)

row = DirectRow(row_key=b"row_key")

row.set_cell("cf1", b"c1", 1)

row.set_cell("cf1", b"c2", 2)

row.set_cell("cf1", b"c3", 3)

row.set_cell("cf1", b"c4", 4)

with self.assertRaises(MaxMutationsError):

mutation_batcher.mutate(row)

@mock.patch("google.cloud.bigtable.batcher.MAX_MUTATIONS", new=3)

def test_mutate_row_with_max_mutations(self):

table = _Table(self.TABLE_NAME)

mutation_batcher = MutationsBatcher(table=table)

row = DirectRow(row_key=b"row_key")

row.set_cell("cf1", b"c1", 1)

row.set_cell("cf1", b"c2", 2)

row.set_cell("cf1", b"c3", 3)

mutation_batcher.mutate(row)

mutation_batcher.flush()

self.assertEqual(table.mutation_calls, 1)

def test_mutate_row_with_max_row_bytes(self):

table = _Table(self.TABLE_NAME)

mutation_batcher = MutationsBatcher(table=table, max_row_bytes=3 * 1024 * 1024)

number_of_bytes = 1 * 1024 * 1024

max_value = b"1" * number_of_bytes

row = DirectRow(row_key=b"row_key")

row.set_cell("cf1", b"c1", max_value)

row.set_cell("cf1", b"c2", max_value)

row.set_cell("cf1", b"c3", max_value)

mutation_batcher.mutate(row)

self.assertEqual(table.mutation_calls, 1)

class _Instance(object):

def __init__(self, client=None):

self._client = client

class _Table(object):

def __init__(self, name, client=None):

self.name = name

self._instance = _Instance(client)

self.mutation_calls = 0

def mutate_rows(self, rows):

self.mutation_calls += 1

return rows