From 49dfe5578c76e2479f40f9c1a90c1d9362321453 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Tim=20Swe=C3=B1a?= Date: Wed, 20 Aug 2025 19:59:20 +0000 Subject: [PATCH] chore: lock global session usage in unit tests --- tests/unit/pandas/io/test_api.py | 86 ++++++++++++++++---------------- tests/unit/test_pandas.py | 15 +++--- 2 files changed, 50 insertions(+), 51 deletions(-) diff --git a/tests/unit/pandas/io/test_api.py b/tests/unit/pandas/io/test_api.py index 1e69fa9df3..eaba412809 100644 --- a/tests/unit/pandas/io/test_api.py +++ b/tests/unit/pandas/io/test_api.py @@ -16,6 +16,7 @@ import pytest +import bigframes.core.global_session import bigframes.dataframe import bigframes.pandas.io.api as bf_io_api import bigframes.session @@ -24,57 +25,56 @@ pytest.importorskip("polars") -@mock.patch( - "bigframes.pandas.io.api._set_default_session_location_if_possible_deferred_query" -) -@mock.patch("bigframes.core.global_session.with_default_session") -def test_read_gbq_colab_dry_run_doesnt_call_set_location( - mock_with_default_session, mock_set_location -): +def test_read_gbq_colab_dry_run_doesnt_call_set_location(): """ Ensure that we don't bind to a location too early. If it's a dry run, the user might not be done typing. """ - mock_df = mock.create_autospec(bigframes.dataframe.DataFrame) - mock_with_default_session.return_value = mock_df + with mock.patch( + "bigframes.pandas.io.api._set_default_session_location_if_possible_deferred_query" + ) as mock_set_location, mock.patch( + "bigframes.core.global_session.with_default_session" + ) as mock_with_default_session, bigframes.core.global_session._global_session_lock: + mock_df = mock.create_autospec(bigframes.dataframe.DataFrame) + mock_with_default_session.return_value = mock_df - query_or_table = "SELECT {param1} AS param1" - sample_pyformat_args = {"param1": "value1"} - bf_io_api._read_gbq_colab( - query_or_table, pyformat_args=sample_pyformat_args, dry_run=True - ) + query_or_table = "SELECT {param1} AS param1" + sample_pyformat_args = {"param1": "value1"} + bf_io_api._read_gbq_colab( + query_or_table, pyformat_args=sample_pyformat_args, dry_run=True + ) - mock_set_location.assert_not_called() + mock_with_default_session.assert_not_called() + mock_set_location.assert_not_called() -@mock.patch( - "bigframes.pandas.io.api._set_default_session_location_if_possible_deferred_query" -) -@mock.patch("bigframes.core.global_session.with_default_session") -def test_read_gbq_colab_calls_set_location( - mock_with_default_session, mock_set_location -): - # Configure the mock for with_default_session to return a DataFrame mock - mock_df = mock.create_autospec(bigframes.dataframe.DataFrame) - mock_with_default_session.return_value = mock_df +def test_read_gbq_colab_calls_set_location(): + with mock.patch( + "bigframes.pandas.io.api._set_default_session_location_if_possible_deferred_query" + ) as mock_set_location, mock.patch( + "bigframes.core.global_session.with_default_session" + ) as mock_with_default_session, bigframes.core.global_session._global_session_lock: + # Configure the mock for with_default_session to return a DataFrame mock + mock_df = mock.create_autospec(bigframes.dataframe.DataFrame) + mock_with_default_session.return_value = mock_df - query_or_table = "SELECT {param1} AS param1" - sample_pyformat_args = {"param1": "value1"} - result = bf_io_api._read_gbq_colab( - query_or_table, pyformat_args=sample_pyformat_args, dry_run=False - ) + query_or_table = "SELECT {param1} AS param1" + sample_pyformat_args = {"param1": "value1"} + result = bf_io_api._read_gbq_colab( + query_or_table, pyformat_args=sample_pyformat_args, dry_run=False + ) - # Make sure that we format the SQL first to prevent syntax errors. - formatted_query = "SELECT 'value1' AS param1" - mock_set_location.assert_called_once() - args, _ = mock_set_location.call_args - assert formatted_query == args[0]() - mock_with_default_session.assert_called_once() + # Make sure that we format the SQL first to prevent syntax errors. + formatted_query = "SELECT 'value1' AS param1" + mock_set_location.assert_called_once() + args, _ = mock_set_location.call_args + assert formatted_query == args[0]() + mock_with_default_session.assert_called_once() - # Check the actual arguments passed to with_default_session - args, kwargs = mock_with_default_session.call_args - assert args[0] == bigframes.session.Session._read_gbq_colab - assert args[1] == query_or_table - assert kwargs["pyformat_args"] == sample_pyformat_args - assert not kwargs["dry_run"] - assert isinstance(result, bigframes.dataframe.DataFrame) + # Check the actual arguments passed to with_default_session + args, kwargs = mock_with_default_session.call_args + assert args[0] == bigframes.session.Session._read_gbq_colab + assert args[1] == query_or_table + assert kwargs["pyformat_args"] == sample_pyformat_args + assert not kwargs["dry_run"] + assert isinstance(result, bigframes.dataframe.DataFrame) diff --git a/tests/unit/test_pandas.py b/tests/unit/test_pandas.py index e8383512a6..8e09e65381 100644 --- a/tests/unit/test_pandas.py +++ b/tests/unit/test_pandas.py @@ -20,9 +20,8 @@ import pandas as pd import pytest -import bigframes.core.global_session import bigframes.dataframe -import bigframes.pandas as bpd +import bigframes.pandas import bigframes.session leading_whitespace = re.compile(r"^\s+", flags=re.MULTILINE) @@ -165,9 +164,9 @@ def test_cut_raises_with_invalid_bins(bins: int, error_message: str): def test_pandas_attribute(): - assert bpd.NA is pd.NA - assert bpd.BooleanDtype is pd.BooleanDtype - assert bpd.Float64Dtype is pd.Float64Dtype - assert bpd.Int64Dtype is pd.Int64Dtype - assert bpd.StringDtype is pd.StringDtype - assert bpd.ArrowDtype is pd.ArrowDtype + assert bigframes.pandas.NA is pd.NA + assert bigframes.pandas.BooleanDtype is pd.BooleanDtype + assert bigframes.pandas.Float64Dtype is pd.Float64Dtype + assert bigframes.pandas.Int64Dtype is pd.Int64Dtype + assert bigframes.pandas.StringDtype is pd.StringDtype + assert bigframes.pandas.ArrowDtype is pd.ArrowDtype