Skip to content
This repository was archived by the owner on Apr 1, 2026. It is now read-only.

Commit 91aa321

Browse files
tseavertseaver
authored
tests: refactor unittests using pytest idioms (#456)
1 parent a535f99 commit 91aa321

15 files changed

Lines changed: 9465 additions & 9625 deletions

‎tests/unit/test_app_profile.py‎

Lines changed: 605 additions & 631 deletions
Large diffs are not rendered by default.

‎tests/unit/test_backup.py‎

Lines changed: 857 additions & 845 deletions
Large diffs are not rendered by default.

‎tests/unit/test_batcher.py‎

Lines changed: 85 additions & 104 deletions
Original file line numberDiff line numberDiff line change
@@ -13,154 +13,135 @@
1313
# limitations under the License.
1414

1515

16-
import unittest
17-
1816
import mock
17+
import pytest
1918

20-
from ._testing import _make_credentials
21-
22-
from google.cloud.bigtable.batcher import MutationsBatcher
2319
from google.cloud.bigtable.row import DirectRow
2420

21+
TABLE_ID = "table-id"
22+
TABLE_NAME = "/tables/" + TABLE_ID
2523

26-
class TestMutationsBatcher(unittest.TestCase):
27-
from grpc import StatusCode
2824

29-
TABLE_ID = "table-id"
30-
TABLE_NAME = "/tables/" + TABLE_ID
25+
def _make_mutation_batcher(table, **kw):
26+
from google.cloud.bigtable.batcher import MutationsBatcher
3127

32-
# RPC Status Codes
33-
SUCCESS = StatusCode.OK.value[0]
28+
return MutationsBatcher(table, **kw)
3429

35-
@staticmethod
36-
def _get_target_class():
37-
from google.cloud.bigtable.table import Table
3830

39-
return Table
31+
def test_mutation_batcher_constructor():
32+
table = _Table(TABLE_NAME)
4033

41-
def _make_table(self, *args, **kwargs):
42-
return self._get_target_class()(*args, **kwargs)
34+
mutation_batcher = _make_mutation_batcher(table)
35+
assert table is mutation_batcher.table
4336

44-
@staticmethod
45-
def _get_target_client_class():
46-
from google.cloud.bigtable.client import Client
4737

48-
return Client
38+
def test_mutation_batcher_mutate_row():
39+
table = _Table(TABLE_NAME)
40+
mutation_batcher = _make_mutation_batcher(table=table)
4941

50-
def _make_client(self, *args, **kwargs):
51-
return self._get_target_client_class()(*args, **kwargs)
42+
rows = [
43+
DirectRow(row_key=b"row_key"),
44+
DirectRow(row_key=b"row_key_2"),
45+
DirectRow(row_key=b"row_key_3"),
46+
DirectRow(row_key=b"row_key_4"),
47+
]
5248

53-
def test_constructor(self):
54-
credentials = _make_credentials()
55-
client = self._make_client(
56-
project="project-id", credentials=credentials, admin=True
57-
)
49+
mutation_batcher.mutate_rows(rows)
50+
mutation_batcher.flush()
5851

59-
instance = client.instance(instance_id="instance-id")
60-
table = self._make_table(self.TABLE_ID, instance)
52+
assert table.mutation_calls == 1
6153

62-
mutation_batcher = MutationsBatcher(table)
63-
self.assertEqual(table, mutation_batcher.table)
6454

65-
def test_mutate_row(self):
66-
table = _Table(self.TABLE_NAME)
67-
mutation_batcher = MutationsBatcher(table=table)
55+
def test_mutation_batcher_mutate():
56+
table = _Table(TABLE_NAME)
57+
mutation_batcher = _make_mutation_batcher(table=table)
6858

69-
rows = [
70-
DirectRow(row_key=b"row_key"),
71-
DirectRow(row_key=b"row_key_2"),
72-
DirectRow(row_key=b"row_key_3"),
73-
DirectRow(row_key=b"row_key_4"),
74-
]
59+
row = DirectRow(row_key=b"row_key")
60+
row.set_cell("cf1", b"c1", 1)
61+
row.set_cell("cf1", b"c2", 2)
62+
row.set_cell("cf1", b"c3", 3)
63+
row.set_cell("cf1", b"c4", 4)
7564

76-
mutation_batcher.mutate_rows(rows)
77-
mutation_batcher.flush()
65+
mutation_batcher.mutate(row)
7866

79-
self.assertEqual(table.mutation_calls, 1)
67+
mutation_batcher.flush()
8068

81-
def test_mutate_rows(self):
82-
table = _Table(self.TABLE_NAME)
83-
mutation_batcher = MutationsBatcher(table=table)
69+
assert table.mutation_calls == 1
8470

85-
row = DirectRow(row_key=b"row_key")
86-
row.set_cell("cf1", b"c1", 1)
87-
row.set_cell("cf1", b"c2", 2)
88-
row.set_cell("cf1", b"c3", 3)
89-
row.set_cell("cf1", b"c4", 4)
9071

91-
mutation_batcher.mutate(row)
72+
def test_mutation_batcher_flush_w_no_rows():
73+
table = _Table(TABLE_NAME)
74+
mutation_batcher = _make_mutation_batcher(table=table)
75+
mutation_batcher.flush()
9276

93-
mutation_batcher.flush()
77+
assert table.mutation_calls == 0
9478

95-
self.assertEqual(table.mutation_calls, 1)
9679

97-
def test_flush_with_no_rows(self):
98-
table = _Table(self.TABLE_NAME)
99-
mutation_batcher = MutationsBatcher(table=table)
100-
mutation_batcher.flush()
80+
def test_mutation_batcher_mutate_w_max_flush_count():
81+
table = _Table(TABLE_NAME)
82+
mutation_batcher = _make_mutation_batcher(table=table, flush_count=3)
10183

102-
self.assertEqual(table.mutation_calls, 0)
84+
row_1 = DirectRow(row_key=b"row_key_1")
85+
row_2 = DirectRow(row_key=b"row_key_2")
86+
row_3 = DirectRow(row_key=b"row_key_3")
10387

104-
def test_add_row_with_max_flush_count(self):
105-
table = _Table(self.TABLE_NAME)
106-
mutation_batcher = MutationsBatcher(table=table, flush_count=3)
88+
mutation_batcher.mutate(row_1)
89+
mutation_batcher.mutate(row_2)
90+
mutation_batcher.mutate(row_3)
10791

108-
row_1 = DirectRow(row_key=b"row_key_1")
109-
row_2 = DirectRow(row_key=b"row_key_2")
110-
row_3 = DirectRow(row_key=b"row_key_3")
92+
assert table.mutation_calls == 1
11193

112-
mutation_batcher.mutate(row_1)
113-
mutation_batcher.mutate(row_2)
114-
mutation_batcher.mutate(row_3)
11594

116-
self.assertEqual(table.mutation_calls, 1)
95+
@mock.patch("google.cloud.bigtable.batcher.MAX_MUTATIONS", new=3)
96+
def test_mutation_batcher_mutate_with_max_mutations_failure():
97+
from google.cloud.bigtable.batcher import MaxMutationsError
11798

118-
@mock.patch("google.cloud.bigtable.batcher.MAX_MUTATIONS", new=3)
119-
def test_mutate_row_with_max_mutations_failure(self):
120-
from google.cloud.bigtable.batcher import MaxMutationsError
99+
table = _Table(TABLE_NAME)
100+
mutation_batcher = _make_mutation_batcher(table=table)
121101

122-
table = _Table(self.TABLE_NAME)
123-
mutation_batcher = MutationsBatcher(table=table)
102+
row = DirectRow(row_key=b"row_key")
103+
row.set_cell("cf1", b"c1", 1)
104+
row.set_cell("cf1", b"c2", 2)
105+
row.set_cell("cf1", b"c3", 3)
106+
row.set_cell("cf1", b"c4", 4)
124107

125-
row = DirectRow(row_key=b"row_key")
126-
row.set_cell("cf1", b"c1", 1)
127-
row.set_cell("cf1", b"c2", 2)
128-
row.set_cell("cf1", b"c3", 3)
129-
row.set_cell("cf1", b"c4", 4)
108+
with pytest.raises(MaxMutationsError):
109+
mutation_batcher.mutate(row)
130110

131-
with self.assertRaises(MaxMutationsError):
132-
mutation_batcher.mutate(row)
133111

134-
@mock.patch("google.cloud.bigtable.batcher.MAX_MUTATIONS", new=3)
135-
def test_mutate_row_with_max_mutations(self):
136-
table = _Table(self.TABLE_NAME)
137-
mutation_batcher = MutationsBatcher(table=table)
112+
@mock.patch("google.cloud.bigtable.batcher.MAX_MUTATIONS", new=3)
113+
def test_mutation_batcher_mutate_w_max_mutations():
114+
table = _Table(TABLE_NAME)
115+
mutation_batcher = _make_mutation_batcher(table=table)
138116

139-
row = DirectRow(row_key=b"row_key")
140-
row.set_cell("cf1", b"c1", 1)
141-
row.set_cell("cf1", b"c2", 2)
142-
row.set_cell("cf1", b"c3", 3)
117+
row = DirectRow(row_key=b"row_key")
118+
row.set_cell("cf1", b"c1", 1)
119+
row.set_cell("cf1", b"c2", 2)
120+
row.set_cell("cf1", b"c3", 3)
143121

144-
mutation_batcher.mutate(row)
145-
mutation_batcher.flush()
122+
mutation_batcher.mutate(row)
123+
mutation_batcher.flush()
146124

147-
self.assertEqual(table.mutation_calls, 1)
125+
assert table.mutation_calls == 1
148126

149-
def test_mutate_row_with_max_row_bytes(self):
150-
table = _Table(self.TABLE_NAME)
151-
mutation_batcher = MutationsBatcher(table=table, max_row_bytes=3 * 1024 * 1024)
152127

153-
number_of_bytes = 1 * 1024 * 1024
154-
max_value = b"1" * number_of_bytes
128+
def test_mutation_batcher_mutate_w_max_row_bytes():
129+
table = _Table(TABLE_NAME)
130+
mutation_batcher = _make_mutation_batcher(
131+
table=table, max_row_bytes=3 * 1024 * 1024
132+
)
155133

156-
row = DirectRow(row_key=b"row_key")
157-
row.set_cell("cf1", b"c1", max_value)
158-
row.set_cell("cf1", b"c2", max_value)
159-
row.set_cell("cf1", b"c3", max_value)
134+
number_of_bytes = 1 * 1024 * 1024
135+
max_value = b"1" * number_of_bytes
160136

161-
mutation_batcher.mutate(row)
137+
row = DirectRow(row_key=b"row_key")
138+
row.set_cell("cf1", b"c1", max_value)
139+
row.set_cell("cf1", b"c2", max_value)
140+
row.set_cell("cf1", b"c3", max_value)
141+
142+
mutation_batcher.mutate(row)
162143

163-
self.assertEqual(table.mutation_calls, 1)
144+
assert table.mutation_calls == 1
164145

165146

166147
class _Instance(object):

0 commit comments

Comments
 (0)