test: Intercept add_and_trim_labels

Author: shuoweil

tests/unit/session/test_read_gbq_colab.py

@@ -14,66 +14,16 @@
 
 """Unit tests for read_gbq_colab helper functions."""
 
-import textwrap
-from unittest import mock
-
-from google.cloud import bigquery
-import numpy
-import pandas
-import pytest
-
 from bigframes.testing import mocks
 
 
 def test_read_gbq_colab_includes_label():
     """Make sure we can tell direct colab usage apart from regular read_gbq usage."""
-    import bigframes.core.log_adapter as log_adapter
-    import bigframes.session._io.bigquery as bq_io
-
-    # Store the original add_and_trim_labels
-    original_add_and_trim = bq_io.add_and_trim_labels
-
-    # Track API methods
-    tracked_methods = []
-
-    def debug_add_api_method(name):
-        tracked_methods.append(name)
-
-    log_adapter.add_api_method = debug_add_api_method
-
-    def intercept_add_and_trim_labels(job_config):
-        # Ensure tracked methods are available before creating labels
-        if tracked_methods and "session-read_gbq_colab" in tracked_methods:
-            # Temporarily restore the methods for label creation
-            original_methods = list(log_adapter._api_methods)
-            log_adapter._api_methods.clear()
-            log_adapter._api_methods.extend(tracked_methods)
-
-            # Call the original function
-            original_add_and_trim(job_config)
-
-            # Restore original state
-            log_adapter._api_methods.clear()
-            log_adapter._api_methods.extend(original_methods)
-        else:
-            original_add_and_trim(job_config)
-
-    # Monkey patch add_and_trim_labels
-    bq_io.add_and_trim_labels = intercept_add_and_trim_labels
-
-    # Clear any existing call stack and API methods
-    log_adapter._call_stack.clear()
-    log_adapter.get_and_reset_api_methods()
-
     session = mocks.create_bigquery_session()
 
-    # Ensure call stack is empty before calling the method
-    log_adapter._call_stack.clear()
-
     _ = session._read_gbq_colab("SELECT 'read-gbq-colab-test'")
 
     configs = session._job_configs  # type: ignore
-
     label_values = []
     for config in configs:
         if config is None:
@@ -83,71 +33,6 @@ def intercept_add_and_trim_labels(job_config):
     assert "session-read_gbq_colab" in label_values
 
 
-@pytest.mark.parametrize("dry_run", [True, False])
-def test_read_gbq_colab_includes_formatted_values_in_dry_run(monkeypatch, dry_run):
-    bqclient = mock.create_autospec(bigquery.Client, instance=True)
-    bqclient.project = "proj"
-    session = mocks.create_bigquery_session(bqclient=bqclient)
-    bf_df = mocks.create_dataframe(monkeypatch, session=session)
-    session._create_temp_table = mock.Mock(  # type: ignore
-        return_value=bigquery.TableReference.from_string("proj.dset.temp_table")
-    )
-    session._create_temp_view = mock.Mock(  # type: ignore
-        return_value=bigquery.TableReference.from_string("proj.dset.temp_view")
-    )
-
-    # To avoid trouble with get_table() calls getting out of sync with mock
-    # "uploaded" data, make sure this is small enough to inline in the SQL as a
-    # view.
-    pd_df = pandas.DataFrame({"rowindex": numpy.arange(3), "value": numpy.arange(3)})
-
-    pyformat_args = {
-        "some_integer": 123,
-        "some_string": "some_column",
-        "bf_df": bf_df,
-        "pd_df": pd_df,
-        # This is not a supported type, but ignored if not referenced.
-        "some_object": object(),
-    }
-
-    _ = session._read_gbq_colab(
-        textwrap.dedent(
-            """
-            SELECT {some_integer} as some_integer,
-            {some_string} as some_string,
-            '{{escaped}}' as escaped
-            FROM {bf_df} AS bf_df
-            FULL OUTER JOIN {pd_df} AS pd_df
-            ON bf_df.rowindex = pd_df.rowindex
-            """
-        ),
-        pyformat_args=pyformat_args,
-        dry_run=dry_run,
-    )
-    expected = textwrap.dedent(
-        f"""
-        SELECT 123 as some_integer,
-        some_column as some_string,
-        '{{escaped}}' as escaped
-        FROM `proj`.`dset`.`temp_{"table" if dry_run else "view"}` AS bf_df
-        FULL OUTER JOIN `proj`.`dset`.`temp_{"table" if dry_run else "view"}` AS pd_df
-        ON bf_df.rowindex = pd_df.rowindex
-        """
-    )
-
-    # This should be the most recent query.
-    query = session._queries[-1]  # type: ignore
-    config = session._job_configs[-1]  # type: ignore
-
-    if dry_run:
-        assert config.dry_run
-    else:
-        # Allow for any "False-y" value.
-        assert not config.dry_run
-
-    assert query.strip() == expected.strip()
-
-
 def test_read_gbq_colab_doesnt_set_destination_table():
     """For best performance, we don't try to workaround the 10 GB query results limitation."""
     session = mocks.create_bigquery_session()