test(bigquery): add unit tests for timeout logic in _pandas_helpers.py

Author: chalmerlowe

tests/unit/job/test_query_pandas.py

@@ -1026,3 +1026,35 @@ def test_query_job_to_geodataframe_delegation(wait_for_query):
         timeout=None,
     )
     assert df is row_iterator.to_geodataframe.return_value
+
+
+@pytest.mark.skipif(pandas is None, reason="Requires `pandas`")
+@mock.patch("google.cloud.bigquery.job.query.wait_for_query")
+def test_query_job_to_dataframe_delegation(wait_for_query):
+    job = _make_job()
+    bqstorage_client = object()
+    timeout = 123.45
+
+    job.to_dataframe(bqstorage_client=bqstorage_client, timeout=timeout)
+
+    wait_for_query.assert_called_once_with(job, None, max_results=None)
+    row_iterator = wait_for_query.return_value
+    row_iterator.to_dataframe.assert_called_once()
+    call_args = row_iterator.to_dataframe.call_args
+    assert call_args.kwargs["timeout"] == timeout
+
+
+@pytest.mark.skipif(pyarrow is None, reason="Requires `pyarrow`")
+@mock.patch("google.cloud.bigquery.job.query.wait_for_query")
+def test_query_job_to_arrow_delegation(wait_for_query):
+    job = _make_job()
+    bqstorage_client = object()
+    timeout = 123.45
+
+    job.to_arrow(bqstorage_client=bqstorage_client, timeout=timeout)
+
+    wait_for_query.assert_called_once_with(job, None, max_results=None)
+    row_iterator = wait_for_query.return_value
+    row_iterator.to_arrow.assert_called_once()
+    call_args = row_iterator.to_arrow.call_args
+    assert call_args.kwargs["timeout"] == timeout

tests/unit/test__pandas_helpers.py

@@ -13,12 +13,14 @@
 # limitations under the License.
 
 import collections
+import concurrent.futures
 import datetime
 import decimal
 import functools
 import gc
 import operator
 import queue
+import time
 from typing import Union
 from unittest import mock
 import warnings
@@ -2177,3 +2179,69 @@ def test_determine_requested_streams_invalid_max_stream_count():
     """Tests that a ValueError is raised if max_stream_count is negative."""
     with pytest.raises(ValueError):
         determine_requested_streams(preserve_order=False, max_stream_count=-1)
+
+
+def test__download_table_bqstorage_w_timeout_error(module_under_test):
+    from google.cloud.bigquery import dataset
+    from google.cloud.bigquery import table
+
+    bqstorage_client = mock.create_autospec(
+        bigquery_storage.BigQueryReadClient, instance=True
+    )
+    # Give it one stream
+    fake_session = mock.Mock(streams=["stream/s0"])
+    bqstorage_client.create_read_session.return_value = fake_session
+
+    table_ref = table.TableReference(
+        dataset.DatasetReference("project-x", "dataset-y"),
+        "table-z",
+    )
+
+    def slow_download_stream(
+        download_state, bqstorage_client, session, stream, worker_queue, page_to_item
+    ):
+        # Block until the main thread sets done=True (which it will on timeout)
+        while not download_state.done:
+            time.sleep(0.01)
+
+    with mock.patch.object(
+        module_under_test, "_download_table_bqstorage_stream", new=slow_download_stream
+    ):
+        # Use a very small timeout
+        result_gen = module_under_test._download_table_bqstorage(
+            "some-project", table_ref, bqstorage_client, timeout=0.01
+        )
+        with pytest.raises(concurrent.futures.TimeoutError, match="timed out"):
+            list(result_gen)
+
+
+def test__download_table_bqstorage_w_timeout_success(module_under_test):
+    from google.cloud.bigquery import dataset
+    from google.cloud.bigquery import table
+
+    bqstorage_client = mock.create_autospec(
+        bigquery_storage.BigQueryReadClient, instance=True
+    )
+    fake_session = mock.Mock(streams=["stream/s0"])
+    bqstorage_client.create_read_session.return_value = fake_session
+
+    table_ref = table.TableReference(
+        dataset.DatasetReference("project-x", "dataset-y"),
+        "table-z",
+    )
+
+    def fast_download_stream(
+        download_state, bqstorage_client, session, stream, worker_queue, page_to_item
+    ):
+        worker_queue.put("result_page")
+
+    with mock.patch.object(
+        module_under_test, "_download_table_bqstorage_stream", new=fast_download_stream
+    ):
+        # Use a generous timeout
+        result_gen = module_under_test._download_table_bqstorage(
+            "some-project", table_ref, bqstorage_client, timeout=10.0
+        )
+        results = list(result_gen)
+
+    assert results == ["result_page"]

tests/unit/test_table.py

@@ -2495,6 +2495,18 @@ def test_to_geodataframe(self):
         else:
             assert not hasattr(df, "crs")
 
+    def test_methods_w_timeout(self):
+        # Ensure that timeout parameter is accepted by all methods and ignored.
+        row_iterator = self._make_one()
+        timeout = 42.0
+
+        # Just calling them to ensure no TypeError is raised
+        row_iterator.to_arrow(timeout=timeout)
+        row_iterator.to_arrow_iterable(timeout=timeout)
+        row_iterator.to_dataframe(timeout=timeout)
+        row_iterator.to_dataframe_iterable(timeout=timeout)
+        row_iterator.to_geodataframe(timeout=timeout)
+
 
 class TestRowIterator(unittest.TestCase):
     PYARROW_MINIMUM_VERSION = str(_versions_helpers._MIN_PYARROW_VERSION)