chore(tests): clean up sample test tables (#1030)

samples/beam/hello_world_write_test.py

@@ -14,45 +14,33 @@
 import os
 import uuid
 
-from google.cloud import bigtable
 import pytest
 
-import hello_world_write
+from . import hello_world_write
+from ..utils import create_table_cm
 
 PROJECT = os.environ["GOOGLE_CLOUD_PROJECT"]
 BIGTABLE_INSTANCE = os.environ["BIGTABLE_INSTANCE"]
-TABLE_ID_PREFIX = "mobile-time-series-{}"
+TABLE_ID = f"mobile-time-series-beam-{str(uuid.uuid4())[:16]}"
 
 
 @pytest.fixture(scope="module", autouse=True)
-def table_id():
-    client = bigtable.Client(project=PROJECT, admin=True)
-    instance = client.instance(BIGTABLE_INSTANCE)
+def table():
+    with create_table_cm(
+        PROJECT, BIGTABLE_INSTANCE, TABLE_ID, {"stats_summary": None}
+    ) as table:
+        yield table
 
-    table_id = TABLE_ID_PREFIX.format(str(uuid.uuid4())[:16])
-    table = instance.table(table_id)
-    if table.exists():
-        table.delete()
 
-    table.create(column_families={"stats_summary": None})
-    yield table_id
-
-    table.delete()
-
-
-def test_hello_world_write(table_id):
+def test_hello_world_write(table):
     hello_world_write.run(
         [
             "--bigtable-project=%s" % PROJECT,
             "--bigtable-instance=%s" % BIGTABLE_INSTANCE,
-            "--bigtable-table=%s" % table_id,
+            "--bigtable-table=%s" % TABLE_ID,
         ]
     )
 
-    client = bigtable.Client(project=PROJECT, admin=True)
-    instance = client.instance(BIGTABLE_INSTANCE)
-    table = instance.table(table_id)
-
     rows = table.read_rows()
     count = 0
     for _ in rows:

samples/hello/async_main.py

@@ -26,14 +26,14 @@
 
 import argparse
 import asyncio
+from ..utils import wait_for_table
 
 # [START bigtable_async_hw_imports]
 from google.cloud import bigtable
 from google.cloud.bigtable.data import row_filters
 from google.cloud.bigtable.data import RowMutationEntry
 from google.cloud.bigtable.data import SetCell
 from google.cloud.bigtable.data import ReadRowsQuery
-
 # [END bigtable_async_hw_imports]
 
 
@@ -65,63 +65,66 @@ async def main(project_id, instance_id, table_id):
         print("Table {} already exists.".format(table_id))
     # [END bigtable_async_hw_create_table]
 
-    # [START bigtable_async_hw_write_rows]
-    print("Writing some greetings to the table.")
-    greetings = ["Hello World!", "Hello Cloud Bigtable!", "Hello Python!"]
-    mutations = []
-    column = "greeting"
-    for i, value in enumerate(greetings):
-        # Note: This example uses sequential numeric IDs for simplicity,
-        # but this can result in poor performance in a production
-        # application.  Since rows are stored in sorted order by key,
-        # sequential keys can result in poor distribution of operations
-        # across nodes.
-        #
-        # For more information about how to design a Bigtable schema for
-        # the best performance, see the documentation:
-        #
-        #     https://cloud.google.com/bigtable/docs/schema-design
-        row_key = "greeting{}".format(i).encode()
-        row_mutation = RowMutationEntry(
-            row_key, SetCell(column_family_id, column, value)
-        )
-        mutations.append(row_mutation)
-    await table.bulk_mutate_rows(mutations)
-    # [END bigtable_async_hw_write_rows]
-
-    # [START bigtable_async_hw_create_filter]
-    # Create a filter to only retrieve the most recent version of the cell
-    # for each column across entire row.
-    row_filter = row_filters.CellsColumnLimitFilter(1)
-    # [END bigtable_async_hw_create_filter]
-
-    # [START bigtable_async_hw_get_with_filter]
-    # [START bigtable_async_hw_get_by_key]
-    print("Getting a single greeting by row key.")
-    key = "greeting0".encode()
-
-    row = await table.read_row(key, row_filter=row_filter)
-    cell = row.cells[0]
-    print(cell.value.decode("utf-8"))
-    # [END bigtable_async_hw_get_by_key]
-    # [END bigtable_async_hw_get_with_filter]
-
-    # [START bigtable_async_hw_scan_with_filter]
-    # [START bigtable_async_hw_scan_all]
-    print("Scanning for all greetings:")
-    query = ReadRowsQuery(row_filter=row_filter)
-    async for row in await table.read_rows_stream(query):
+    try:
+        # let table creation complete
+        wait_for_table(admin_table)
+        # [START bigtable_async_hw_write_rows]
+        print("Writing some greetings to the table.")
+        greetings = ["Hello World!", "Hello Cloud Bigtable!", "Hello Python!"]
+        mutations = []
+        column = "greeting"
+        for i, value in enumerate(greetings):
+            # Note: This example uses sequential numeric IDs for simplicity,
+            # but this can result in poor performance in a production
+            # application.  Since rows are stored in sorted order by key,
+            # sequential keys can result in poor distribution of operations
+            # across nodes.
+            #
+            # For more information about how to design a Bigtable schema for
+            # the best performance, see the documentation:
+            #
+            #     https://cloud.google.com/bigtable/docs/schema-design
+            row_key = "greeting{}".format(i).encode()
+            row_mutation = RowMutationEntry(
+                row_key, SetCell(column_family_id, column, value)
+            )
+            mutations.append(row_mutation)
+        await table.bulk_mutate_rows(mutations)
+        # [END bigtable_async_hw_write_rows]
+
+        # [START bigtable_async_hw_create_filter]
+        # Create a filter to only retrieve the most recent version of the cell
+        # for each column across entire row.
+        row_filter = row_filters.CellsColumnLimitFilter(1)
+        # [END bigtable_async_hw_create_filter]
+
+        # [START bigtable_async_hw_get_with_filter]
+        # [START bigtable_async_hw_get_by_key]
+        print("Getting a single greeting by row key.")
+        key = "greeting0".encode()
+
+        row = await table.read_row(key, row_filter=row_filter)
         cell = row.cells[0]
         print(cell.value.decode("utf-8"))
-    # [END bigtable_async_hw_scan_all]
-    # [END bigtable_async_hw_scan_with_filter]
-
-    # [START bigtable_async_hw_delete_table]
-    # the async client only supports the data API. Table deletion as an admin operation
-    # use admin client to create the table
-    print("Deleting the {} table.".format(table_id))
-    admin_table.delete()
-    # [END bigtable_async_hw_delete_table]
+        # [END bigtable_async_hw_get_by_key]
+        # [END bigtable_async_hw_get_with_filter]
+
+        # [START bigtable_async_hw_scan_with_filter]
+        # [START bigtable_async_hw_scan_all]
+        print("Scanning for all greetings:")
+        query = ReadRowsQuery(row_filter=row_filter)
+        async for row in await table.read_rows_stream(query):
+            cell = row.cells[0]
+            print(cell.value.decode("utf-8"))
+        # [END bigtable_async_hw_scan_all]
+        # [END bigtable_async_hw_scan_with_filter]
+    finally:
+        # [START bigtable_async_hw_delete_table]
+        # the async client only supports the data API. Table deletion as an admin operation
+        # use admin client to create the table
+        print("Deleting the {} table.".format(table_id))
+        admin_table.delete()
+        # [END bigtable_async_hw_delete_table]
 
 
 if __name__ == "__main__":

samples/hello/async_main_test.py

@@ -13,27 +13,24 @@
 # limitations under the License.
 
 import os
-import random
 import asyncio
+import uuid
 
-from async_main import main
+from .async_main import main
 
 PROJECT = os.environ["GOOGLE_CLOUD_PROJECT"]
 BIGTABLE_INSTANCE = os.environ["BIGTABLE_INSTANCE"]
-TABLE_NAME_FORMAT = "hello-world-test-{}"
-TABLE_NAME_RANGE = 10000
+TABLE_ID = f"hello-world-test-async-{str(uuid.uuid4())[:16]}"
 
 
 def test_async_main(capsys):
-    table_name = TABLE_NAME_FORMAT.format(random.randrange(TABLE_NAME_RANGE))
-
-    asyncio.run(main(PROJECT, BIGTABLE_INSTANCE, table_name))
+    asyncio.run(main(PROJECT, BIGTABLE_INSTANCE, TABLE_ID))
 
     out, _ = capsys.readouterr()
-    assert "Creating the {} table.".format(table_name) in out
+    assert "Creating the {} table.".format(TABLE_ID) in out
     assert "Writing some greetings to the table." in out
     assert "Getting a single greeting by row key." in out
     assert "Hello World!" in out
     assert "Scanning for all greetings" in out
     assert "Hello Cloud Bigtable!" in out
-    assert "Deleting the {} table.".format(table_name) in out
+    assert "Deleting the {} table.".format(TABLE_ID) in out

samples/hello/main.py

@@ -25,6 +25,7 @@
 """
 
 import argparse
+from ..utils import wait_for_table
 
 # [START bigtable_hw_imports]
 import datetime
@@ -60,63 +61,68 @@ def main(project_id, instance_id, table_id):
         print("Table {} already exists.".format(table_id))
     # [END bigtable_hw_create_table]
 
-    # [START bigtable_hw_write_rows]
-    print("Writing some greetings to the table.")
-    greetings = ["Hello World!", "Hello Cloud Bigtable!", "Hello Python!"]
-    rows = []
-    column = "greeting".encode()
-    for i, value in enumerate(greetings):
-        # Note: This example uses sequential numeric IDs for simplicity,
-        # but this can result in poor performance in a production
-        # application.  Since rows are stored in sorted order by key,
-        # sequential keys can result in poor distribution of operations
-        # across nodes.
-        #
-        # For more information about how to design a Bigtable schema for
-        # the best performance, see the documentation:
-        #
-        #     https://cloud.google.com/bigtable/docs/schema-design
-        row_key = "greeting{}".format(i).encode()
-        row = table.direct_row(row_key)
-        row.set_cell(
-            column_family_id, column, value, timestamp=datetime.datetime.utcnow()
-        )
-        rows.append(row)
-    table.mutate_rows(rows)
-    # [END bigtable_hw_write_rows]
-
-    # [START bigtable_hw_create_filter]
-    # Create a filter to only retrieve the most recent version of the cell
-    # for each column across entire row.
-    row_filter = row_filters.CellsColumnLimitFilter(1)
-    # [END bigtable_hw_create_filter]
-
-    # [START bigtable_hw_get_with_filter]
-    # [START bigtable_hw_get_by_key]
-    print("Getting a single greeting by row key.")
-    key = "greeting0".encode()
-
-    row = table.read_row(key, row_filter)
-    cell = row.cells[column_family_id][column][0]
-    print(cell.value.decode("utf-8"))
-    # [END bigtable_hw_get_by_key]
-    # [END bigtable_hw_get_with_filter]
-
-    # [START bigtable_hw_scan_with_filter]
-    # [START bigtable_hw_scan_all]
-    print("Scanning for all greetings:")
-    partial_rows = table.read_rows(filter_=row_filter)
-
-    for row in partial_rows:
+    try:
+        # let table creation complete
+        wait_for_table(table)
+
+        # [START bigtable_hw_write_rows]
+        print("Writing some greetings to the table.")
+        greetings = ["Hello World!", "Hello Cloud Bigtable!", "Hello Python!"]
+        rows = []
+        column = "greeting".encode()
+        for i, value in enumerate(greetings):
+            # Note: This example uses sequential numeric IDs for simplicity,
+            # but this can result in poor performance in a production
+            # application.  Since rows are stored in sorted order by key,
+            # sequential keys can result in poor distribution of operations
+            # across nodes.
+            #
+            # For more information about how to design a Bigtable schema for
+            # the best performance, see the documentation:
+            #
+            #     https://cloud.google.com/bigtable/docs/schema-design
+            row_key = "greeting{}".format(i).encode()
+            row = table.direct_row(row_key)
+            row.set_cell(
+                column_family_id, column, value, timestamp=datetime.datetime.utcnow()
+            )
+            rows.append(row)
+        table.mutate_rows(rows)
+        # [END bigtable_hw_write_rows]
+
+        # [START bigtable_hw_create_filter]
+        # Create a filter to only retrieve the most recent version of the cell
+        # for each column across entire row.
+        row_filter = row_filters.CellsColumnLimitFilter(1)
+        # [END bigtable_hw_create_filter]
+
+        # [START bigtable_hw_get_with_filter]
+        # [START bigtable_hw_get_by_key]
+        print("Getting a single greeting by row key.")
+        key = "greeting0".encode()
+
+        row = table.read_row(key, row_filter)
         cell = row.cells[column_family_id][column][0]
         print(cell.value.decode("utf-8"))
-    # [END bigtable_hw_scan_all]
-    # [END bigtable_hw_scan_with_filter]
-
-    # [START bigtable_hw_delete_table]
-    print("Deleting the {} table.".format(table_id))
-    table.delete()
-    # [END bigtable_hw_delete_table]
+        # [END bigtable_hw_get_by_key]
+        # [END bigtable_hw_get_with_filter]
+
+        # [START bigtable_hw_scan_with_filter]
+        # [START bigtable_hw_scan_all]
+        print("Scanning for all greetings:")
+        partial_rows = table.read_rows(filter_=row_filter)
+
+        for row in partial_rows:
+            cell = row.cells[column_family_id][column][0]
+            print(cell.value.decode("utf-8"))
+        # [END bigtable_hw_scan_all]
+        # [END bigtable_hw_scan_with_filter]
+
+    finally:
+        # [START bigtable_hw_delete_table]
+        print("Deleting the {} table.".format(table_id))
+        table.delete()
+        # [END bigtable_hw_delete_table]
 
 
 if __name__ == "__main__":

samples/hello/main_test.py

@@ -13,26 +13,23 @@
 # limitations under the License.
 
 import os
-import random
+import uuid
 
-from main import main
+from .main import main
 
 PROJECT = os.environ["GOOGLE_CLOUD_PROJECT"]
 BIGTABLE_INSTANCE = os.environ["BIGTABLE_INSTANCE"]
-TABLE_NAME_FORMAT = "hello-world-test-{}"
-TABLE_NAME_RANGE = 10000
+TABLE_ID = f"hello-world-test-{str(uuid.uuid4())[:16]}"
 
 
 def test_main(capsys):
-    table_name = TABLE_NAME_FORMAT.format(random.randrange(TABLE_NAME_RANGE))
-
-    main(PROJECT, BIGTABLE_INSTANCE, table_name)
+    main(PROJECT, BIGTABLE_INSTANCE, TABLE_ID)
 
     out, _ = capsys.readouterr()
-    assert "Creating the {} table.".format(table_name) in out
+    assert "Creating the {} table.".format(TABLE_ID) in out
     assert "Writing some greetings to the table." in out
     assert "Getting a single greeting by row key." in out
     assert "Hello World!" in out
     assert "Scanning for all greetings" in out
     assert "Hello Cloud Bigtable!" in out
-    assert "Deleting the {} table.".format(table_name) in out
+    assert "Deleting the {} table.".format(TABLE_ID) in out