Skip to content
This repository was archived by the owner on Mar 6, 2026. It is now read-only.

Commit ef809a0

Browse files
aribrayaribray
committed
add tqdm mock patch
1 parent 09a6608 commit ef809a0

File tree

1 file changed

+34
-24
lines changed

1 file changed

+34
-24
lines changed

‎tests/unit/job/test_query_pandas.py‎

Lines changed: 34 additions & 24 deletions
Original file line numberDiff line numberDiff line change
@@ -12,6 +12,7 @@
1212
# See the License for the specific language governing permissions and
1313
# limitations under the License.
1414

15+
import concurrent.futures
1516
import copy
1617
import json
1718

@@ -36,7 +37,7 @@
3637
except (ImportError, AttributeError): # pragma: NO COVER
3738
geopandas = None
3839
try:
39-
from tqdm import tqdm
40+
import tqdm
4041
except (ImportError, AttributeError): # pragma: NO COVER
4142
tqdm = None
4243

@@ -300,7 +301,8 @@ def test_to_arrow_max_results_no_progress_bar():
300301

301302

302303
@pytest.mark.skipif(tqdm is None, reason="Requires `tqdm`")
303-
def test_to_arrow_w_tqdm_w_query_plan():
304+
@mock.patch("tqdm.tqdm")
305+
def test_to_arrow_w_tqdm_w_query_plan(tqdm_mock):
304306
from google.cloud.bigquery import table
305307
from google.cloud.bigquery.job import QueryJob as target_class
306308
from google.cloud.bigquery.schema import SchemaField
@@ -338,23 +340,26 @@ def test_to_arrow_w_tqdm_w_query_plan():
338340
result_patch = mock.patch(
339341
"google.cloud.bigquery.job.QueryJob.result",
340342
side_effect=[
343+
concurrent.futures.TimeoutError,
344+
concurrent.futures.TimeoutError,
341345
row_iterator,
342346
],
343347
)
344348

345-
with result_patch as result_patch_tqdm, reload_patch:
349+
with result_patch as tqdm_mock, reload_patch:
346350
tbl = job.to_arrow(progress_bar_type="tqdm", create_bqstorage_client=False)
347351

348-
assert result_patch_tqdm.call_count == 1
352+
assert tqdm_mock.call_count == 3
349353
assert isinstance(tbl, pyarrow.Table)
350354
assert tbl.num_rows == 2
351-
result_patch_tqdm.assert_called_with(
355+
tqdm_mock.assert_called_with(
352356
timeout=_PROGRESS_BAR_UPDATE_INTERVAL, max_results=None
353357
)
354358

355359

356360
@pytest.mark.skipif(tqdm is None, reason="Requires `tqdm`")
357-
def test_to_arrow_w_tqdm_w_pending_status():
361+
@mock.patch("tqdm.tqdm")
362+
def test_to_arrow_w_tqdm_w_pending_status(tqdm_mock):
358363
from google.cloud.bigquery import table
359364
from google.cloud.bigquery.job import QueryJob as target_class
360365
from google.cloud.bigquery.schema import SchemaField
@@ -391,16 +396,16 @@ def test_to_arrow_w_tqdm_w_pending_status():
391396
)
392397
result_patch = mock.patch(
393398
"google.cloud.bigquery.job.QueryJob.result",
394-
side_effect=[row_iterator],
399+
side_effect=[concurrent.futures.TimeoutError, row_iterator],
395400
)
396401

397-
with result_patch as result_patch_tqdm, reload_patch:
402+
with result_patch as tqdm_mock, reload_patch:
398403
tbl = job.to_arrow(progress_bar_type="tqdm", create_bqstorage_client=False)
399404

400-
assert result_patch_tqdm.call_count == 1
405+
assert tqdm_mock.call_count == 2
401406
assert isinstance(tbl, pyarrow.Table)
402407
assert tbl.num_rows == 2
403-
result_patch_tqdm.assert_called_with(
408+
tqdm_mock.assert_called_with(
404409
timeout=_PROGRESS_BAR_UPDATE_INTERVAL, max_results=None
405410
)
406411

@@ -748,7 +753,8 @@ def test_to_dataframe_with_progress_bar(tqdm_mock):
748753

749754

750755
@pytest.mark.skipif(tqdm is None, reason="Requires `tqdm`")
751-
def test_to_dataframe_w_tqdm_pending():
756+
@mock.patch("tqdm.tqdm")
757+
def test_to_dataframe_w_tqdm_pending(tqdm_mock):
752758
from google.cloud.bigquery import table
753759
from google.cloud.bigquery.job import QueryJob as target_class
754760
from google.cloud.bigquery.schema import SchemaField
@@ -787,23 +793,24 @@ def test_to_dataframe_w_tqdm_pending():
787793
)
788794
result_patch = mock.patch(
789795
"google.cloud.bigquery.job.QueryJob.result",
790-
side_effect=[row_iterator],
796+
side_effect=[concurrent.futures.TimeoutError, row_iterator],
791797
)
792798

793-
with result_patch as result_patch_tqdm, reload_patch:
799+
with result_patch as tqdm_mock, reload_patch:
794800
df = job.to_dataframe(progress_bar_type="tqdm", create_bqstorage_client=False)
795801

796-
assert result_patch_tqdm.call_count == 1
802+
assert tqdm_mock.call_count == 2
797803
assert isinstance(df, pandas.DataFrame)
798804
assert len(df) == 4 # verify the number of rows
799805
assert list(df) == ["name", "age"] # verify the column names
800-
result_patch_tqdm.assert_called_with(
806+
tqdm_mock.assert_called_with(
801807
timeout=_PROGRESS_BAR_UPDATE_INTERVAL, max_results=None
802808
)
803809

804810

805811
@pytest.mark.skipif(tqdm is None, reason="Requires `tqdm`")
806-
def test_to_dataframe_w_tqdm():
812+
@mock.patch("tqdm.tqdm")
813+
def test_to_dataframe_w_tqdm(tqdm_mock):
807814
from google.cloud.bigquery import table
808815
from google.cloud.bigquery.job import QueryJob as target_class
809816
from google.cloud.bigquery.schema import SchemaField
@@ -843,24 +850,27 @@ def test_to_dataframe_w_tqdm():
843850
result_patch = mock.patch(
844851
"google.cloud.bigquery.job.QueryJob.result",
845852
side_effect=[
853+
concurrent.futures.TimeoutError,
854+
concurrent.futures.TimeoutError,
846855
row_iterator,
847856
],
848857
)
849858

850-
with result_patch as result_patch_tqdm, reload_patch:
859+
with result_patch as tqdm_mock, reload_patch:
851860
df = job.to_dataframe(progress_bar_type="tqdm", create_bqstorage_client=False)
852861

853-
assert result_patch_tqdm.call_count == 1
862+
assert tqdm_mock.call_count == 3
854863
assert isinstance(df, pandas.DataFrame)
855864
assert len(df) == 4 # verify the number of rows
856865
assert list(df), ["name", "age"] # verify the column names
857-
result_patch_tqdm.assert_called_with(
866+
tqdm_mock.assert_called_with(
858867
timeout=_PROGRESS_BAR_UPDATE_INTERVAL, max_results=None
859868
)
860869

861870

862871
@pytest.mark.skipif(tqdm is None, reason="Requires `tqdm`")
863-
def test_to_dataframe_w_tqdm_max_results():
872+
@mock.patch("tqdm.tqdm")
873+
def test_to_dataframe_w_tqdm_max_results(tqdm_mock):
864874
from google.cloud.bigquery import table
865875
from google.cloud.bigquery.job import QueryJob as target_class
866876
from google.cloud.bigquery.schema import SchemaField
@@ -894,16 +904,16 @@ def test_to_dataframe_w_tqdm_max_results():
894904
)
895905
result_patch = mock.patch(
896906
"google.cloud.bigquery.job.QueryJob.result",
897-
side_effect=[row_iterator],
907+
side_effect=[concurrent.futures.TimeoutError, row_iterator],
898908
)
899909

900-
with result_patch as result_patch_tqdm, reload_patch:
910+
with result_patch as tqdm_mock, reload_patch:
901911
job.to_dataframe(
902912
progress_bar_type="tqdm", create_bqstorage_client=False, max_results=3
903913
)
904914

905-
assert result_patch_tqdm.call_count == 1
906-
result_patch_tqdm.assert_called_with(
915+
assert tqdm_mock.call_count == 2
916+
tqdm_mock.assert_called_with(
907917
timeout=_PROGRESS_BAR_UPDATE_INTERVAL, max_results=3
908918
)
909919

0 commit comments

Comments
 (0)