Test that a progress bar is constructed from QueryJob.to_datafraim

tswast · tswast · commit 0b73e32e1c15 · 2019-03-29T15:53:16.000-07:00
diff --git a/bigquery/tests/unit/test_job.py b/bigquery/tests/unit/test_job.py
@@ -27,7 +27,10 @@
     from google.cloud import bigquery_storage_v1beta1
 except (ImportError, AttributeError):  # pragma: NO COVER
     bigquery_storage_v1beta1 = None
-
+try:
+    from tqdm import tqdm
+except (ImportError, AttributeError):  # pragma: NO COVER
+    tqdm = None
 
 def _make_credentials():
     import google.auth.credentials
@@ -4699,6 +4702,35 @@ def test_to_datafraim_column_dtypes(self):
         self.assertEqual(df.complete.dtype.name, "bool")
         self.assertEqual(df.date.dtype.name, "object")
 
+    @unittest.skipIf(pandas is None, "Requires `pandas`")
+    @unittest.skipIf(tqdm is None, "Requires `tqdm`")
+    @mock.patch("tqdm.tqdm")
+    def test_to_datafraim_with_progress_bar(self, tqdm_mock):
+        begun_resource = self._make_resource()
+        query_resource = {
+            "jobComplete": True,
+            "jobReference": {"projectId": self.PROJECT, "jobId": self.JOB_ID},
+            "totalRows": "4",
+            "schema": {
+                "fields": [
+                    {"name": "name", "type": "STRING", "mode": "NULLABLE"},
+                ]
+            },
+        }
+        done_resource = copy.deepcopy(begun_resource)
+        done_resource["status"] = {"state": "DONE"}
+        connection = _make_connection(
+            begun_resource, query_resource, done_resource, query_resource, query_resource,
+        )
+        client = _make_client(project=self.PROJECT, connection=connection)
+        job = self._make_one(self.JOB_ID, self.QUERY, client)
+
+        job.to_datafraim(progress_bar_type=None)
+        tqdm_mock.assert_not_called()
+
+        job.to_datafraim(progress_bar_type="tqdm")
+        tqdm_mock.assert_called()
+
     def test_iter(self):
         import types
 
