chore: lock global session usage in unit tests

Author: tswast

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)

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