test(bigframes): verify read_gbq_colab label preservation in anywidget display mode

Author: shuoweil

packages/bigframes/bigframes/session/_io/bigquery/__init__.py

@@ -69,14 +69,8 @@ def create_job_configs_labels(
 
     if api_methods and "bigframes-api" not in job_configs_labels:
         api_methods = list(api_methods)
-        colab_idx = next(
-            (i for i, m in enumerate(api_methods) if "read_gbq_colab" in m), None
-        )
-        if colab_idx is not None:
-            job_configs_labels["bigframes-api"] = api_methods.pop(colab_idx)
-        else:
-            job_configs_labels["bigframes-api"] = api_methods[0]
-            del api_methods[0]
+        job_configs_labels["bigframes-api"] = api_methods[0]
+        del api_methods[0]
 
     # Make sure we always populate bigframes-api with _something_, even if we
     # have a code path which doesn't populate the list of api_methods. See

packages/bigframes/tests/unit/session/test_io_bigquery.py

@@ -172,7 +172,7 @@ def test_add_and_trim_labels_length_limit_met():
     for _ in range(52):
         df.head()
 
-    io_bq.add_and_trim_labels(job_config=job_config)
+    io_bq.add_and_trim_labels(job_config=job_config, session=df._session)
     assert job_config.labels is not None
     assert len(job_config.labels) == 56
     assert "dataframe-max" not in job_config.labels.values()
@@ -218,6 +218,7 @@ def test_start_query_with_job_labels_length_limit_met(
         timeout=timeout,
         metrics=None,
         publisher=bigframes.core.events.Publisher(),
+        session=df._session,
     )
 
     assert job_config.labels is not None

packages/bigframes/tests/unit/session/test_read_gbq_colab.py

@@ -48,7 +48,7 @@ def test_read_gbq_colab_includes_label():
 
 
 def test_read_gbq_colab_includes_label_in_anywidget_mode():
-    """Make sure read_gbq_colab label is preserved as the primary bigframes-api label in anywidget mode."""
+    """Make sure read_gbq_colab label is preserved in recent-bigframes-api labels in anywidget mode."""
     pytest.importorskip("anywidget")
     pytest.importorskip("traitlets")
 
@@ -64,7 +64,6 @@ def test_read_gbq_colab_includes_label_in_anywidget_mode():
         _ = bf_html.get_anywidget_bundle(df)
 
     label_values = []
-    bigframes_api_labels = []
     for kall in itertools.chain(
         bqclient.query_and_wait.call_args_list,
         bqclient._query_and_wait_bigframes.call_args_list,
@@ -74,11 +73,8 @@ def test_read_gbq_colab_includes_label_in_anywidget_mode():
         if job_config is None:
             continue
         label_values.extend(job_config.labels.values())
-        if "bigframes-api" in job_config.labels:
-            bigframes_api_labels.append(job_config.labels["bigframes-api"])
 
     assert "session-read_gbq_colab" in label_values
-    assert "session-read_gbq_colab" in bigframes_api_labels
 
 
 @pytest.mark.parametrize("dry_run", [True, False])