test(bigquery-magics): make table_id parsing check version-agnostic (#17562)

chalmerlowe · web-flow · commit 17bef417db64 · 2026-06-25T13:11:20.000-04:00
&gt; [!NOTE]
&gt; These changes are at a high level unrelated. They are being submitted
as a single corrective action. Separating them leads to a chicken or the
egg set of failures (i.e. no one PR will pass all CI/CD tests with out
the other changes).

---

### test(bigquery-magics): make table_id parsing check version-agnostic

The unit test `test_bigquery_magic_query_variable_not_identifier` was
asserting that a specific error message ("must be a fully-qualified ID")
was in the stderr output.

However, in a recent versions of dependencies, this error message
changed (depending on whether tests were run under Python 3.10 or
3.11+). There is a key term that is common to both of the error
messages: `table_id`.

This change confirms that we get one of the two expected error messages
by checking for `table_id`, making it compatible regardless of which
runtime the tests run under.

---

### refactor(magics): replace fragile psutil socket check with client
close mock spy

The system test `test_bigquery_magic` was using `psutil` to count open
sockets before and after the test to detect leaks (i.e. ensure all the
system resources were released).

This approach proved to be fragile across different environments (e.g.,
Kokoro).

This change replaces the socket counting with a mock spy on
`google.cloud.bigquery.Client.close` to verify that the client is
explicitly closed, which is the intended mechanism for releasing system
resources.

---

### test(bigquery-magics): simplify IPython initialization in
deprecation test

The unit test `test_cell_magic_engine_bigframes_warning` in
`test_deprecation.py` was conditionally initializing IPython. Only one
of the code paths was being exercised so we were getting a coverage
issue.

This change simplifies the initialization by unconditionally calling
`start_ipython()` and assigning the returned instance to `ip` ensuring a
clean and consistent state for the test and avoiding the need for a
conditional check.
diff --git a/packages/bigquery-magics/tests/system/test_bigquery.py b/packages/bigquery-magics/tests/system/test_bigquery.py
@@ -15,18 +15,17 @@
 """System tests for Jupyter/IPython connector."""
 
 import re
+from unittest import mock
 
+import google.cloud.bigquery
+import pandas
 from IPython.testing import globalipapp
 from IPython.utils import io
-import pandas
-import psutil
 
 
 def test_bigquery_magic():
     globalipapp.start_ipython()
     ip = globalipapp.get_ipython()
-    current_process = psutil.Process()
-    conn_count_start = len(current_process.net_connections())
 
     ip.extension_manager.load_extension("bigquery_magics")
     sql = """
@@ -40,10 +39,17 @@ def test_bigquery_magic():
         ORDER BY view_count DESC
         LIMIT 10
     """
-    with io.capture_output() as captured:
-        result = ip.run_cell_magic("bigquery", "--use_rest_api", sql)
-
-    conn_count_end = len(current_process.net_connections())
+    with mock.patch.object(
+        google.cloud.bigquery.Client,
+        "close",
+        autospec=True,
+        side_effect=google.cloud.bigquery.Client.close,
+    ) as mock_close:
+        with io.capture_output() as captured:
+            result = ip.run_cell_magic("bigquery", "--use_rest_api", sql)
+
+        # Verify that client close is explicitly called to release sockets.
+        mock_close.assert_called_once()
 
     lines = re.split("\n|\r", captured.stdout)
     # Removes blanks & terminal code (result of display clearing)
@@ -53,8 +59,3 @@ def test_bigquery_magic():
     assert isinstance(result, pandas.DataFrame)
     assert len(result) == 10  # verify row count
     assert list(result) == ["url", "view_count"]  # verify column names
-
-    # NOTE: For some reason, the number of open sockets is sometimes one *less*
-    # than expected when running system tests on Kokoro, thus using the <= assertion.
-    # That's still fine, however, since the sockets are apparently not leaked.
-    assert conn_count_end <= conn_count_start  # system resources are released
diff --git a/packages/bigquery-magics/tests/unit/bigquery/test_bigquery.py b/packages/bigquery-magics/tests/unit/bigquery/test_bigquery.py
@@ -2974,7 +2974,7 @@ def test_bigquery_magic_query_variable_not_identifier():
     # considered a table name, thus we expect an error that the table ID is not valid.
     output = captured_io.stderr
     assert "ERROR:" in output
-    assert "must be a fully-qualified ID" in output
+    assert "table_id" in output
 
 
 @pytest.mark.usefixtures("mock_credentials")
diff --git a/packages/bigquery-magics/tests/unit/bigquery/test_deprecation.py b/packages/bigquery-magics/tests/unit/bigquery/test_deprecation.py
@@ -54,13 +54,10 @@ def test_query_with_bigframes_warning(mock_ipython):
 def test_cell_magic_engine_bigframes_warning(mock_ipython):
     from unittest import mock
 
-    from IPython.testing.globalipapp import get_ipython
+    from IPython.testing.globalipapp import get_ipython, start_ipython
 
+    start_ipython()
     ip = get_ipython()
-    if ip is None:
-        from IPython.testing.globalipapp import start_ipython
-
-        ip = start_ipython()
 
     ip.extension_manager.load_extension("bigquery_magics")
 
