moved table manegement into shared utils module

samples/beam/hello_world_write_test.py

@@ -13,30 +13,22 @@
 # limitations under the License.
 import os
 
-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 = "mobile-time-series-beam"
 
 
 @pytest.fixture(scope="module", autouse=True)
-def table_id():
-    client = bigtable.Client(project=PROJECT, admin=True)
-    instance = client.instance(BIGTABLE_INSTANCE)
-
-    table_id = TABLE_ID
-    table = instance.table(table_id)
-    if table.exists():
-        table.delete()
-
-    table.create(column_families={"stats_summary": None})
-    yield table
-
-    table.delete()
+def table():
+    with create_table_cm(
+        PROJECT, BIGTABLE_INSTANCE, TABLE_ID, {"stats_summary": None}
+    ) as table:
+        yield table
 
 
 def test_hello_world_write(table):

samples/beam/noxfile.py

@@ -22,7 +22,6 @@
 
 import nox
 
-
 # WARNING - WARNING - WARNING - WARNING - WARNING
 # WARNING - WARNING - WARNING - WARNING - WARNING
 #           DO NOT EDIT THIS FILE EVER!
@@ -158,6 +157,7 @@ def blacken(session: nox.sessions.Session) -> None:
 # format = isort + black
 #
 
+
 @nox.session
 def format(session: nox.sessions.Session) -> None:
     """
@@ -185,7 +185,9 @@ def _session_tests(
     session: nox.sessions.Session, post_install: Callable = None
 ) -> None:
     # check for presence of tests
-    test_list = glob.glob("**/*_test.py", recursive=True) + glob.glob("**/test_*.py", recursive=True)
+    test_list = glob.glob("**/*_test.py", recursive=True) + glob.glob(
+        "**/test_*.py", recursive=True
+    )
     test_list.extend(glob.glob("**/tests", recursive=True))
 
     if len(test_list) == 0:
@@ -207,9 +209,7 @@ def _session_tests(
 
     if os.path.exists("requirements-test.txt"):
         if os.path.exists("constraints-test.txt"):
-            session.install(
-                "-r", "requirements-test.txt", "-c", "constraints-test.txt"
-            )
+            session.install("-r", "requirements-test.txt", "-c", "constraints-test.txt")
         else:
             session.install("-r", "requirements-test.txt")
         with open("requirements-test.txt") as rtfile:
@@ -222,9 +222,9 @@ def _session_tests(
         post_install(session)
 
     if "pytest-parallel" in packages:
-        concurrent_args.extend(['--workers', 'auto', '--tests-per-worker', 'auto'])
+        concurrent_args.extend(["--workers", "auto", "--tests-per-worker", "auto"])
     elif "pytest-xdist" in packages:
-        concurrent_args.extend(['-n', 'auto'])
+        concurrent_args.extend(["-n", "auto"])
 
     session.run(
         "pytest",
@@ -254,7 +254,7 @@ def py(session: nox.sessions.Session) -> None:
 
 
 def _get_repo_root() -> Optional[str]:
-    """ Returns the root folder of the project. """
+    """Returns the root folder of the project."""
     # Get root of this repository. Assume we don't have directories nested deeper than 10 items.
     p = Path(os.getcwd())
     for i in range(10):

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
-from google.api_core.exceptions import PreconditionFailed
 # [END bigtable_async_hw_imports]
 
 
@@ -64,74 +64,67 @@ async def main(project_id, instance_id, table_id):
     else:
         print("Table {} already exists.".format(table_id))
     # [END bigtable_async_hw_create_table]
-    # let table creation complete
-    attempts = 0
-    table_ready = False
-    while not table_ready and attempts < 10:
-        try:
-            table_ready = table.exists()
-        except PreconditionFailed:
-            print("Waiting for table to become ready...")
-        attempts += 1
-        await asyncio.sleep(5)
-
-    # [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

@@ -14,31 +14,22 @@
 
 import os
 import asyncio
-from google.cloud import bigtable
 
-from async_main import main
+from .async_main import main
 
 PROJECT = os.environ["GOOGLE_CLOUD_PROJECT"]
 BIGTABLE_INSTANCE = os.environ["BIGTABLE_INSTANCE"]
 TABLE_ID = "hello-world-test-async"
 
 
 def test_async_main(capsys):
-    try:
-        asyncio.run(main(PROJECT, BIGTABLE_INSTANCE, TABLE_ID))
+    asyncio.run(main(PROJECT, BIGTABLE_INSTANCE, TABLE_ID))
 
-        out, _ = capsys.readouterr()
-        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_ID) in out
-    finally:
-        # delete table
-        client = bigtable.Client(PROJECT, admin=True)
-        instance = client.instance(BIGTABLE_INSTANCE)
-        table = instance.table(TABLE_ID)
-        if table.exists():
-            table.delete()
+    out, _ = capsys.readouterr()
+    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_ID) in out