WIP: improve load_table_dataframe sample tests.

tswast · tswast · commit b04a3c6bab01 · 2019-08-23T16:48:19.000-07:00
diff --git a/bigquery/samples/load_table_dataframe.py b/bigquery/samples/load_table_dataframe.py
@@ -13,7 +13,7 @@
 # limitations under the License.
 
 
-def main(client, table_id):
+def load_table_dataframe(client, table_id):
     # [START bigquery_load_table_dataframe]
     from google.cloud import bigquery
     import pandas
@@ -38,7 +38,7 @@ def main(client, table_id):
         # table. The schema is used to assist in data type definitions.
         schema=[
             # Specify the type of columns whose type cannot be auto-detected. For
-            # example the "title" column uses pandas dtype "object", so it's
+            # example the "title" column uses pandas dtype "object", so its
             # data type is ambiguous.
             bigquery.SchemaField("title", bigquery.enums.SqlTypeNames.STRING),
             # Indexes are written if included in the schema by name.
@@ -60,4 +60,4 @@ def main(client, table_id):
     import sys
     from google.cloud import bigquery
 
-    main(bigquery.Client(), sys.argv[1])
+    load_table_dataframe(bigquery.Client(), sys.argv[1])
diff --git a/bigquery/samples/tests/test_load_table_dataframe.py b/bigquery/samples/tests/test_load_table_dataframe.py
@@ -18,23 +18,11 @@
 
 
 pytest.importorskip("pandas")
+pytest.importorskip("pyarrow")
 
 
-@pytest.mark.parametrize("parquet_engine", ["pyarrow", "fastparquet"])
-def test_main():
-    if parquet_engine == "pyarrow" and pyarrow is None:
-        pytest.skip("Requires `pyarrow`")
-    if parquet_engine == "fastparquet" and fastparquet is None:
-        pytest.skip("Requires `fastparquet`")
-
-    pandas.set_option("io.parquet.engine", parquet_engine)
-
-    dataset_id = "load_table_from_dataframe_{}".format(_millis())
-    dataset = bigquery.Dataset(client.dataset(dataset_id))
-    client.create_dataset(dataset)
-    to_delete.append(dataset)
-
-    load_table_dataframe.main()
+def test_load_table_dataframe(client, random_table_id):
+    load_table_dataframe.load_table_dataframe(client, random_table_id)
 
     column_names = [field.name for field in table.schema]
     assert sorted(column_names) == ["release_year", "title", "wikidata_id"]
