add test fixtures

Author: Linchin

samples/conftest.py

@@ -12,6 +12,7 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+import datetime
 import os
 
 import pytest
@@ -20,3 +21,37 @@
 @pytest.fixture(scope="session")
 def project_id():
     return os.environ["GOOGLE_CLOUD_PROJECT"]
+
+
+@pytest.fixture(scope="session")
+def dataset_id():
+    return os.environ["GOOGLE_CLOUD_PROJECT"]
+
+def _make_dataset(project_id, bq_client, location):
+    from google.cloud import bigquery
+
+    dataset_name = prefixer.create_prefix()
+
+    dataset_id = "{}.{}".format(project_id, dataset_name)
+    dataset = bigquery.Dataset(dataset_id)
+    dataset.location = location
+    created_dataset = bq_client.create_dataset(dataset)
+
+    return created_dataset
+
+
+@pytest.fixture(scope="session")
+def dataset(project_id):
+    from google.cloud import bigquery
+
+    client = bigquery.Client()
+    dataset_suffix = datetime.datetime.now().strftime("%y%m%d_%H%M%S")
+    dataset_name = "samples_tests_" + dataset_suffix
+
+    dataset_id = "{}.{}".format(project_id, dataset_name)
+    dataset = bigquery.Dataset(dataset_id)
+    dataset.location = "us-east7"
+    created_dataset = client.create_dataset(dataset)
+    yield created_dataset
+
+    client.delete_dataset(created_dataset, delete_contents=True)

samples/pyarrow/append_rows_with_arrow.py

@@ -15,14 +15,14 @@
 # limitations under the License.
 import datetime
 import decimal
-import pandas as pd
-import pyarrow as pa
 
 from google.cloud import bigquery
 from google.cloud.bigquery import enums
-
 from google.cloud.bigquery_storage_v1 import types as gapic_types
 from google.cloud.bigquery_storage_v1.writer import AppendRowsStream
+import pandas as pd
+
+import pyarrow as pa
 
 
 def bqstorage_write_client():
@@ -31,7 +31,7 @@ def bqstorage_write_client():
     return bigquery_storage_v1.BigQueryWriteClient()
 
 
-def make_table(project_id, dataset, table_id, bq_client):
+def make_table(project_id, dataset, bq_client):
     schema = [
         bigquery.SchemaField("bool_col", enums.SqlTypeNames.BOOLEAN),
         bigquery.SchemaField("int64_col", enums.SqlTypeNames.INT64),
@@ -59,6 +59,7 @@ def make_table(project_id, dataset, table_id, bq_client):
             range_element_type="TIMESTAMP",
         ),
     ]
+    table_id = "append_rows_w_arrow_test"
     table_id_full = f"{project_id}.{dataset}.{table_id}"
     bq_table = bigquery.Table(table_id_full, schema=schema)
     created_table = bq_client.create_table(bq_table)
@@ -167,8 +168,8 @@ def append_rows(bqstorage_write_client, table):
         print(e)
 
 
-def main(project_id, dataset_id, table_id):
+def main(project_id, dataset_id):
     write_client = bqstorage_write_client()
     bq_client = bigquery.Client()
-    table = make_table(project_id, dataset_id, table_id, bq_client)
+    table = make_table(project_id, dataset_id, bq_client)
     append_rows(write_client, table)

samples/pyarrow/append_rows_with_arrow_test.py

@@ -15,7 +15,7 @@
 from . import append_rows_with_arrow
 
 
-def test_read_query_results(capsys, project_id, dataset_id, table_id):
-    append_rows_with_arrow.main(project_id, dataset_id, table_id)
+def test_read_query_results(capsys, project_id, dataset_id):
+    append_rows_with_arrow.main(project_id, dataset_id)
     out, _ = capsys.readouterr()
     assert "append_result" in out