feat(bigtable): separate the snippet test

googleapis · May 20, 2020 · 909ba28 · 909ba28
1 parent 57e59b7
commit 909ba28
Show file tree

Hide file tree

Showing 4 changed files with 28 additions and 18 deletions.
diff --git a/docs/snippets_table.py b/docs/snippets_table.py
@@ -620,7 +620,7 @@ def test_bigtable_create_update_delete_column_family():
     assert column_family_id not in column_families
 
 
-def test_bigtable_add_row_add_row_range_add_row_range_from_keys_add_row_range_by_prefix_from_keys():
+def test_bigtable_add_row_add_row_range_add_row_range_from_keys():
     row_keys = [
         b"row_key_1",
         b"row_key_2",
@@ -631,7 +631,6 @@ def test_bigtable_add_row_add_row_range_add_row_range_from_keys_add_row_range_by
         b"row_key_7",
         b"row_key_8",
         b"row_key_9",
-        b"sample_row_key_1",
     ]
 
     rows = []
@@ -692,8 +691,26 @@ def test_bigtable_add_row_add_row_range_add_row_range_from_keys_add_row_range_by
     expected_row_keys = [b"row_key_3", b"row_key_4", b"row_key_5", b"row_key_6"]
     found_row_keys = [row.row_key for row in read_rows]
     assert found_row_keys == expected_row_keys
+    table.truncate(timeout=200)
+
+
+def test_bigtable_add_row_range_with_prefix():
+    row_keys = [
+        b"row_key_1",
+        b"row_key_2",
+        b"row_key_3",
+        b"sample_row_key_1",
+        b"sample_row_key_2",
+    ]
+
+    rows = []
+    for row_key in row_keys:
+        row = Config.TABLE.row(row_key)
+        row.set_cell(COLUMN_FAMILY_ID, COL_NAME1, CELL_VAL1)
+        rows.append(row)
+    Config.TABLE.mutate_rows(rows)
 
-    # [START bigtable_row_range_by_prefix_from_keys]
+    # [START bigtable_add_row_range_with_prefix]
     from google.cloud.bigtable import Client
     from google.cloud.bigtable.row_set import RowSet
 
@@ -702,24 +719,17 @@ def test_bigtable_add_row_add_row_range_add_row_range_from_keys_add_row_range_by
     table = instance.table(TABLE_ID)
 
     row_set = RowSet()
-    row_set.add_row_range_by_prefix_from_keys("row")
-    # [END bigtable_row_range_by_prefix_from_keys]
+    row_set.add_row_range_with_prefix("row")
+    # [END bigtable_add_row_range_with_prefix]
 
     read_rows = table.read_rows(row_set=row_set)
     expected_row_keys = [
         b"row_key_1",
         b"row_key_2",
         b"row_key_3",
-        b"row_key_4",
-        b"row_key_5",
-        b"row_key_6",
-        b"row_key_7",
-        b"row_key_8",
-        b"row_key_9",
     ]
     found_row_keys = [row.row_key for row in read_rows]
     assert found_row_keys == expected_row_keys
-
     table.truncate(timeout=200)
 
 

diff --git a/google/cloud/bigtable/row_set.py b/google/cloud/bigtable/row_set.py
@@ -109,14 +109,14 @@ def add_row_range_from_keys(
         row_range = RowRange(start_key, end_key, start_inclusive, end_inclusive)
         self.row_ranges.append(row_range)
 
-    def add_row_range_by_prefix_from_keys(self, row_key_prefix):
+    def add_row_range_with_prefix(self, row_key_prefix):
         """Add row range to row_ranges list that start with the row_key_prefix from the row keys
 
         For example:
 
         .. literalinclude:: snippets_table.py
-            :start-after: [START bigtable_row_range_by_prefix_from_keys]
-            :end-before: [END bigtable_row_range_by_prefix_from_keys]
+            :start-after: [START bigtable_add_row_range_with_prefix]
+            :end-before: [END bigtable_add_row_range_with_prefix]
 
         :type row_key_prefix: str
         :param row_key_prefix: To retrieve  all rows that start with this row key prefix.

diff --git a/tests/system.py b/tests/system.py
@@ -1032,7 +1032,7 @@ def test_add_row_range_by_prefix_from_keys(self):
         self._table.mutate_rows(rows)
 
         row_set = RowSet()
-        row_set.add_row_range_by_prefix_from_keys("row")
+        row_set.add_row_range_with_prefix("row")
 
         read_rows = self._table.yield_rows(row_set=row_set)
 

diff --git a/tests/unit/test_row_set.py b/tests/unit/test_row_set.py
@@ -167,9 +167,9 @@ def test_add_row_range_from_keys(self):
         )
         self.assertEqual(row_set.row_ranges[0].end_key, b"row_key9")
 
-    def test_add_row_range_by_prefix_from_keys(self):
+    def test_add_row_range_with_prefix(self):
         row_set = self._make_one()
-        row_set.add_row_range_by_prefix_from_keys("row")
+        row_set.add_row_range_with_prefix("row")
         self.assertEqual(row_set.row_ranges[0].end_key, b"rox")
 
     def test__update_message_request(self):