From 0358f4f68e601b5e594f4dece41651ada405ce8e Mon Sep 17 00:00:00 2001 From: chalmer lowe Date: Tue, 23 Jun 2026 16:07:11 -0400 Subject: [PATCH 1/5] test(bigquery-magics): support both old and new table_id parse error formats --- .../bigquery-magics/tests/unit/bigquery/test_bigquery.py | 5 ++++- 1 file changed, 4 insertions(+), 1 deletion(-) diff --git a/packages/bigquery-magics/tests/unit/bigquery/test_bigquery.py b/packages/bigquery-magics/tests/unit/bigquery/test_bigquery.py index 3be93afdb721..eae93684b422 100644 --- a/packages/bigquery-magics/tests/unit/bigquery/test_bigquery.py +++ b/packages/bigquery-magics/tests/unit/bigquery/test_bigquery.py @@ -2974,7 +2974,10 @@ 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 ( + "must be a fully-qualified ID" in output + or "Could not parse table_id." in output + ) @pytest.mark.usefixtures("mock_credentials") From 15bb190211039a1ec898b613e10f6c3a1c5b76f9 Mon Sep 17 00:00:00 2001 From: chalmer lowe Date: Tue, 23 Jun 2026 18:39:38 -0400 Subject: [PATCH 2/5] test(bigquery-magics): make table_id parsing check version-agnostic --- .../bigquery-magics/tests/unit/bigquery/test_bigquery.py | 5 +---- 1 file changed, 1 insertion(+), 4 deletions(-) diff --git a/packages/bigquery-magics/tests/unit/bigquery/test_bigquery.py b/packages/bigquery-magics/tests/unit/bigquery/test_bigquery.py index eae93684b422..55dd8de50119 100644 --- a/packages/bigquery-magics/tests/unit/bigquery/test_bigquery.py +++ b/packages/bigquery-magics/tests/unit/bigquery/test_bigquery.py @@ -2974,10 +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 - or "Could not parse table_id." in output - ) + assert "table_id" in output @pytest.mark.usefixtures("mock_credentials") From 3329de11f48929b86b7f18adfc54d60166ab895f Mon Sep 17 00:00:00 2001 From: chalmer lowe Date: Tue, 23 Jun 2026 18:46:35 -0400 Subject: [PATCH 3/5] test(bigquery-magics): format and fix code coverage for deprecation unit test --- .../bigquery_magics/_versions_helpers.py | 2 +- .../bigquery_magics/bigquery.py | 19 +- .../bigquery-magics/bigquery_magics/config.py | 2 +- .../bigquery-magics/bigquery_magics/core.py | 2 +- .../bigquery_magics/line_arg_parser/lexer.py | 2 +- .../tests/system/test_bigquery.py | 4 +- .../tests/unit/bigquery/test_bigquery.py | 169 +++++++++++------- .../tests/unit/bigquery/test_deprecation.py | 9 +- .../tests/unit/bigquery/test_pyformat.py | 2 +- 9 files changed, 123 insertions(+), 88 deletions(-) diff --git a/packages/bigquery-magics/bigquery_magics/_versions_helpers.py b/packages/bigquery-magics/bigquery_magics/_versions_helpers.py index 192011a5b633..4c7d5337e928 100644 --- a/packages/bigquery-magics/bigquery_magics/_versions_helpers.py +++ b/packages/bigquery-magics/bigquery_magics/_versions_helpers.py @@ -16,8 +16,8 @@ from typing import Any -from google.cloud.bigquery import exceptions import packaging.version +from google.cloud.bigquery import exceptions _MIN_BQ_STORAGE_VERSION = packaging.version.Version("2.0.0") diff --git a/packages/bigquery-magics/bigquery_magics/bigquery.py b/packages/bigquery-magics/bigquery_magics/bigquery.py index d6fb565e4ea2..cd8dbb3eb3e5 100644 --- a/packages/bigquery-magics/bigquery_magics/bigquery.py +++ b/packages/bigquery-magics/bigquery_magics/bigquery.py @@ -106,33 +106,33 @@ from __future__ import print_function import ast -from concurrent import futures import copy import json import re import sys import threading import time -from typing import Any, List, Tuple import warnings +from concurrent import futures +from typing import Any, List, Tuple import IPython # type: ignore -from IPython.core import magic_arguments # type: ignore -from IPython.core.getipython import get_ipython +import pandas from google.api_core.exceptions import NotFound from google.cloud import bigquery from google.cloud.bigquery import exceptions from google.cloud.bigquery.dataset import DatasetReference from google.cloud.bigquery.dbapi import _helpers from google.cloud.bigquery.job import QueryJobConfig -import pandas +from IPython.core import magic_arguments # type: ignore +from IPython.core.getipython import get_ipython -from bigquery_magics import core -from bigquery_magics import line_arg_parser as lap import bigquery_magics._versions_helpers import bigquery_magics.config import bigquery_magics.graph_server as graph_server import bigquery_magics.pyformat +from bigquery_magics import core +from bigquery_magics import line_arg_parser as lap try: from google.cloud import bigquery_storage # type: ignore @@ -471,8 +471,9 @@ def _parse_magic_args(line: str) -> Tuple[List[Any], Any]: except lap.ParseError as exc: raise ValueError( - "Unrecognized input, are option values correct? " - "Error details: {}".format(exc.args[0]) + "Unrecognized input, are option values correct? Error details: {}".format( + exc.args[0] + ) ) from exc params = [] diff --git a/packages/bigquery-magics/bigquery_magics/config.py b/packages/bigquery-magics/bigquery_magics/config.py index 2b5719ca3163..4596ddf61cca 100644 --- a/packages/bigquery-magics/bigquery_magics/config.py +++ b/packages/bigquery-magics/bigquery_magics/config.py @@ -12,9 +12,9 @@ # See the License for the specific language governing permissions and # limitations under the License. +import warnings from dataclasses import dataclass from typing import Optional -import warnings import google.api_core.client_options as client_options import google.cloud.bigquery as bigquery diff --git a/packages/bigquery-magics/bigquery_magics/core.py b/packages/bigquery-magics/bigquery_magics/core.py index 7fc52207ebaf..e9fab22078e9 100644 --- a/packages/bigquery-magics/bigquery_magics/core.py +++ b/packages/bigquery-magics/bigquery_magics/core.py @@ -18,9 +18,9 @@ from google.api_core import client_info from google.cloud import bigquery -from bigquery_magics import environment import bigquery_magics.config import bigquery_magics.version +from bigquery_magics import environment context = bigquery_magics.config.context diff --git a/packages/bigquery-magics/bigquery_magics/line_arg_parser/lexer.py b/packages/bigquery-magics/bigquery_magics/line_arg_parser/lexer.py index 6e8b4cc9637b..180dd8b9528f 100644 --- a/packages/bigquery-magics/bigquery_magics/line_arg_parser/lexer.py +++ b/packages/bigquery-magics/bigquery_magics/line_arg_parser/lexer.py @@ -12,10 +12,10 @@ # See the License for the specific language governing permissions and # limitations under the License. -from collections import OrderedDict, namedtuple import enum import itertools import re +from collections import OrderedDict, namedtuple Token = namedtuple("Token", ("type_", "lexeme", "pos")) StateTransition = namedtuple("StateTransition", ("new_state", "total_offset")) diff --git a/packages/bigquery-magics/tests/system/test_bigquery.py b/packages/bigquery-magics/tests/system/test_bigquery.py index e870a9535ade..d23e7be7efb0 100644 --- a/packages/bigquery-magics/tests/system/test_bigquery.py +++ b/packages/bigquery-magics/tests/system/test_bigquery.py @@ -16,10 +16,10 @@ import re -from IPython.testing import globalipapp -from IPython.utils import io import pandas import psutil +from IPython.testing import globalipapp +from IPython.utils import io def test_bigquery_magic(): diff --git a/packages/bigquery-magics/tests/unit/bigquery/test_bigquery.py b/packages/bigquery-magics/tests/unit/bigquery/test_bigquery.py index 55dd8de50119..2260406faa61 100644 --- a/packages/bigquery-magics/tests/unit/bigquery/test_bigquery.py +++ b/packages/bigquery-magics/tests/unit/bigquery/test_bigquery.py @@ -12,7 +12,6 @@ # See the License for the specific language governing permissions and # limitations under the License. -from concurrent import futures import contextlib import copy import json @@ -21,22 +20,23 @@ import re import sys import tempfile -from unittest import mock import warnings +from concurrent import futures +from unittest import mock +import google.auth.credentials +import google.cloud.bigquery._http +import google.cloud.bigquery.exceptions import IPython -from IPython.testing import globalipapp import IPython.utils.io as io +import pandas +import pytest from google.api_core import exceptions -import google.auth.credentials from google.cloud import bigquery from google.cloud.bigquery import exceptions as bq_exceptions from google.cloud.bigquery import job, table -import google.cloud.bigquery._http -import google.cloud.bigquery.exceptions from google.cloud.bigquery.retry import DEFAULT_TIMEOUT -import pandas -import pytest +from IPython.testing import globalipapp import bigquery_magics import bigquery_magics.bigquery as magics @@ -366,8 +366,9 @@ def test__make_bqstorage_client_true_obsolete_dependency(): "google-cloud-bigquery-storage is outdated" ), ) - with patcher, pytest.raises( - google.cloud.bigquery.exceptions.LegacyBigQueryStorageError + with ( + patcher, + pytest.raises(google.cloud.bigquery.exceptions.LegacyBigQueryStorageError), ): magics._make_bqstorage_client(test_client, {}) @@ -509,9 +510,11 @@ def test_bigquery_graph_spanner_graph_notebook_missing(monkeypatch): ) query_job_mock.to_dataframe.return_value = result - with run_query_patch as run_query_mock, ( - bqstorage_client_patch - ), display_patch as display_mock: + with ( + run_query_patch as run_query_mock, + bqstorage_client_patch, + display_patch as display_mock, + ): run_query_mock.return_value = query_job_mock return_value = ip.run_cell_magic("bigquery", "--graph", sql) @@ -564,9 +567,11 @@ def test_bigquery_graph_int_result(monkeypatch): ) query_job_mock.to_dataframe.return_value = result - with run_query_patch as run_query_mock, ( - bqstorage_client_patch - ), display_patch as display_mock: + with ( + run_query_patch as run_query_mock, + bqstorage_client_patch, + display_patch as display_mock, + ): run_query_mock.return_value = query_job_mock return_value = ip.run_cell_magic("bigquery", "--graph", sql) @@ -619,9 +624,11 @@ def test_bigquery_graph_str_result(monkeypatch): ) query_job_mock.to_dataframe.return_value = result - with run_query_patch as run_query_mock, ( - bqstorage_client_patch - ), display_patch as display_mock: + with ( + run_query_patch as run_query_mock, + bqstorage_client_patch, + display_patch as display_mock, + ): run_query_mock.return_value = query_job_mock return_value = ip.run_cell_magic("bigquery", "--graph", sql) @@ -692,9 +699,11 @@ def test_bigquery_graph_json_json_result(monkeypatch): query_job_mock.configuration.destination.dataset_id = DATASET_ID query_job_mock.configuration.destination.table_id = TABLE_ID - with run_query_patch as run_query_mock, ( - bqstorage_client_patch - ), display_patch as display_mock: + with ( + run_query_patch as run_query_mock, + bqstorage_client_patch, + display_patch as display_mock, + ): run_query_mock.return_value = query_job_mock try: return_value = ip.run_cell_magic("bigquery", "--graph", sql) @@ -761,9 +770,11 @@ def test_bigquery_graph_json_result(monkeypatch): query_job_mock.configuration.destination.dataset_id = DATASET_ID query_job_mock.configuration.destination.table_id = TABLE_ID - with run_query_patch as run_query_mock, ( - bqstorage_client_patch - ), display_patch as display_mock: + with ( + run_query_patch as run_query_mock, + bqstorage_client_patch, + display_patch as display_mock, + ): run_query_mock.return_value = query_job_mock return_value = ip.run_cell_magic("bigquery", "--graph", sql) @@ -871,9 +882,11 @@ def test_bigquery_graph_size_exceeds_max(monkeypatch): query_job_mock.configuration.destination.dataset_id = DATASET_ID query_job_mock.configuration.destination.table_id = TABLE_ID - with run_query_patch as run_query_mock, ( - bqstorage_client_patch - ), display_patch as display_mock: + with ( + run_query_patch as run_query_mock, + bqstorage_client_patch, + display_patch as display_mock, + ): run_query_mock.return_value = query_job_mock ip.run_cell_magic("bigquery", "--graph", sql) @@ -932,9 +945,11 @@ def test_bigquery_graph_size_exceeds_query_result_max(monkeypatch): query_job_mock.configuration.destination.dataset_id = DATASET_ID query_job_mock.configuration.destination.table_id = TABLE_ID - with run_query_patch as run_query_mock, ( - bqstorage_client_patch - ), display_patch as display_mock: + with ( + run_query_patch as run_query_mock, + bqstorage_client_patch, + display_patch as display_mock, + ): run_query_mock.return_value = query_job_mock ip.run_cell_magic("bigquery", "--graph", sql) @@ -994,9 +1009,11 @@ def test_bigquery_graph_with_args_serialization(monkeypatch): query_job_mock.configuration.destination.dataset_id = DATASET_ID query_job_mock.configuration.destination.table_id = TABLE_ID - with run_query_patch as run_query_mock, ( - bqstorage_client_patch - ), display_patch as display_mock: + with ( + run_query_patch as run_query_mock, + bqstorage_client_patch, + display_patch as display_mock, + ): run_query_mock.return_value = query_job_mock endpoint = "https://example.com" @@ -1076,9 +1093,11 @@ def test_bigquery_graph_colab(monkeypatch): query_job_mock.configuration.destination.dataset_id = DATASET_ID query_job_mock.configuration.destination.table_id = "test_destination_table" - with run_query_patch as run_query_mock, ( - bqstorage_client_patch - ), display_patch as display_mock: + with ( + run_query_patch as run_query_mock, + bqstorage_client_patch, + display_patch as display_mock, + ): run_query_mock.return_value = query_job_mock try: return_value = ip.run_cell_magic("bigquery", "--graph", sql) @@ -1205,9 +1224,11 @@ def test_bigquery_graph_missing_spanner_deps(monkeypatch): ) query_job_mock.to_dataframe.return_value = result - with run_query_patch as run_query_mock, ( - bqstorage_client_patch - ), display_patch as display_mock: + with ( + run_query_patch as run_query_mock, + bqstorage_client_patch, + display_patch as display_mock, + ): run_query_mock.return_value = query_job_mock with pytest.raises(ImportError): try: @@ -1485,9 +1506,13 @@ def test_bigquery_magic_default_connection_user_agent_vscode_extension( home_dir_patch = mock.patch("pathlib.Path.home", return_value=user_home) - with conn_patch as conn, ( - run_query_patch - ), default_patch, env_patch, home_dir_patch: + with ( + conn_patch as conn, + run_query_patch, + default_patch, + env_patch, + home_dir_patch, + ): ip.run_cell_magic("bigquery", "", "SELECT 17 as num") expected_user_agents = [ @@ -1564,9 +1589,13 @@ def custom_import_module_side_effect(name, package=None): "importlib.import_module", side_effect=custom_import_module_side_effect ) - with conn_patch as conn, ( - run_query_patch - ), default_patch, env_patch, extension_import_patch: + with ( + conn_patch as conn, + run_query_patch, + default_patch, + env_patch, + extension_import_patch, + ): ip.run_cell_magic("bigquery", "", "SELECT 17 as num") client_info_arg = conn.call_args[1].get("client_info") @@ -1695,9 +1724,11 @@ def test_bigquery_magic_with_bqstorage_from_argument(monkeypatch): google.cloud.bigquery.job.QueryJob, instance=True ) query_job_mock.to_dataframe.return_value = result - with run_query_patch as run_query_mock, ( - bqstorage_client_patch - ), warnings.catch_warnings(record=True) as warned: + with ( + run_query_patch as run_query_mock, + bqstorage_client_patch, + warnings.catch_warnings(record=True) as warned, + ): run_query_mock.return_value = query_job_mock return_value = ip.run_cell_magic("bigquery", "--use_bqstorage_api", sql) @@ -1901,11 +1932,12 @@ def test_bigquery_magic_w_max_results_query_job_results_fails(): ) query_job_mock.result.side_effect = [[], OSError] - with pytest.raises( - OSError - ), client_query_patch as client_query_mock, ( - default_patch - ), close_transports_patch as close_transports: + with ( + pytest.raises(OSError), + client_query_patch as client_query_mock, + default_patch, + close_transports_patch as close_transports, + ): client_query_mock.return_value = query_job_mock ip.run_cell_magic("bigquery", "--max_results=5", sql) @@ -2911,9 +2943,10 @@ def test_bigquery_magic_nonexisting_query_variable(): ip.user_ns.pop("custom_query", None) # Make sure the variable does NOT exist. cell_body = "$custom_query" # Referring to a non-existing variable name. - with pytest.raises( - NameError, match=r".*custom_query does not exist.*" - ), run_query_patch as run_query_mock: + with ( + pytest.raises(NameError, match=r".*custom_query does not exist.*"), + run_query_patch as run_query_mock, + ): ip.run_cell_magic("bigquery", "", cell_body) run_query_mock.assert_not_called() @@ -2928,9 +2961,10 @@ def test_bigquery_magic_empty_query_variable_name(): run_query_patch = mock.patch("bigquery_magics.bigquery._run_query", autospec=True) cell_body = "$" # Not referring to any variable (name omitted). - with pytest.raises( - NameError, match=r"(?i).*missing query variable name.*" - ), run_query_patch as run_query_mock: + with ( + pytest.raises(NameError, match=r"(?i).*missing query variable name.*"), + run_query_patch as run_query_mock, + ): ip.run_cell_magic("bigquery", "", cell_body) run_query_mock.assert_not_called() @@ -2949,9 +2983,10 @@ def test_bigquery_magic_query_variable_non_string(ipython_ns_cleanup): ip.user_ns["custom_query"] = object() cell_body = "$custom_query" # Referring to a non-string variable. - with pytest.raises( - TypeError, match=r".*must be a string or a bytes-like.*" - ), run_query_patch as run_query_mock: + with ( + pytest.raises(TypeError, match=r".*must be a string or a bytes-like.*"), + run_query_patch as run_query_mock, + ): ip.run_cell_magic("bigquery", "", cell_body) run_query_mock.assert_not_called() @@ -3076,9 +3111,11 @@ def test_bigquery_magic_create_dataset_fails(): autospec=True, ) - with pytest.raises( - OSError - ), create_dataset_if_necessary_patch, close_transports_patch as close_transports: + with ( + pytest.raises(OSError), + create_dataset_if_necessary_patch, + close_transports_patch as close_transports, + ): ip.run_cell_magic( "bigquery", "--destination_table dataset_id.table_id", diff --git a/packages/bigquery-magics/tests/unit/bigquery/test_deprecation.py b/packages/bigquery-magics/tests/unit/bigquery/test_deprecation.py index f8fd25b7b625..a4b59bf0eb87 100644 --- a/packages/bigquery-magics/tests/unit/bigquery/test_deprecation.py +++ b/packages/bigquery-magics/tests/unit/bigquery/test_deprecation.py @@ -14,8 +14,8 @@ import pytest -from bigquery_magics import bigquery as magics import bigquery_magics.config +from bigquery_magics import bigquery as magics @pytest.fixture(autouse=True) @@ -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") diff --git a/packages/bigquery-magics/tests/unit/bigquery/test_pyformat.py b/packages/bigquery-magics/tests/unit/bigquery/test_pyformat.py index 3b9e8dc9f5de..97be2ca0afed 100644 --- a/packages/bigquery-magics/tests/unit/bigquery/test_pyformat.py +++ b/packages/bigquery-magics/tests/unit/bigquery/test_pyformat.py @@ -17,8 +17,8 @@ from typing import List from unittest import mock -from IPython.testing import globalipapp import pytest +from IPython.testing import globalipapp @pytest.mark.parametrize( From 80d94f63475a371c4c4d9e33f43adf9bb1874374 Mon Sep 17 00:00:00 2001 From: chalmer lowe Date: Tue, 23 Jun 2026 19:07:32 -0400 Subject: [PATCH 4/5] fix(magics): resolve socket leak in system tests by tracking dynamic credential requests --- .../tests/system/test_bigquery.py | 71 +++++++++++++++---- 1 file changed, 56 insertions(+), 15 deletions(-) diff --git a/packages/bigquery-magics/tests/system/test_bigquery.py b/packages/bigquery-magics/tests/system/test_bigquery.py index d23e7be7efb0..441bc60924b2 100644 --- a/packages/bigquery-magics/tests/system/test_bigquery.py +++ b/packages/bigquery-magics/tests/system/test_bigquery.py @@ -14,7 +14,10 @@ """System tests for Jupyter/IPython connector.""" +import contextlib +import gc import re +import time import pandas import psutil @@ -22,28 +25,66 @@ from IPython.utils import io +@contextlib.contextmanager +def patch_tracked_requests(): + """Context manager to patch google-auth requests and track/close their HTTP sessions. + + This prevents socket leaks in system tests that use Workload Identity or metadata server auth. + """ + import google.auth.transport.requests + + original_init = google.auth.transport.requests.Request.__init__ + tracked_requests = [] + + def patched_init(self, session=None): + original_init(self, session=session) + if session is None: + tracked_requests.append(self) + + google.auth.transport.requests.Request.__init__ = patched_init + try: + yield tracked_requests + finally: + google.auth.transport.requests.Request.__init__ = original_init + for req in tracked_requests: + if hasattr(req, "session") and req.session is not None: + req.session.close() + + def test_bigquery_magic(): globalipapp.start_ipython() ip = globalipapp.get_ipython() current_process = psutil.Process() + + # GC to ensure clean starting state + gc.collect() conn_count_start = len(current_process.net_connections()) - ip.extension_manager.load_extension("bigquery_magics") - sql = """ - SELECT - CONCAT( - 'https://stackoverflow.com/questions/', - CAST(id as STRING)) as url, - view_count - FROM `bigquery-public-data.stackoverflow.posts_questions` - WHERE tags like '%google-bigquery%' - ORDER BY view_count DESC - LIMIT 10 - """ - with io.capture_output() as captured: - result = ip.run_cell_magic("bigquery", "--use_rest_api", sql) + with patch_tracked_requests(): + ip.extension_manager.load_extension("bigquery_magics") + sql = """ + SELECT + CONCAT( + 'https://stackoverflow.com/questions/', + CAST(id as STRING)) as url, + view_count + FROM `bigquery-public-data.stackoverflow.posts_questions` + WHERE tags like '%google-bigquery%' + ORDER BY view_count DESC + LIMIT 10 + """ + with io.capture_output() as captured: + result = ip.run_cell_magic("bigquery", "--use_rest_api", sql) + + # Force garbage collection to sweep unreferenced socket objects + gc.collect() - conn_count_end = len(current_process.net_connections()) + # Wait a bit for the asynchronous channel teardown to complete and the socket to be closed. + for _ in range(30): + conn_count_end = len(current_process.net_connections()) + if conn_count_end <= conn_count_start: + break + time.sleep(0.1) lines = re.split("\n|\r", captured.stdout) # Removes blanks & terminal code (result of display clearing) From 33a862f802c7dc7370be7285d6584f6c5f5b4ab6 Mon Sep 17 00:00:00 2001 From: chalmer lowe Date: Tue, 23 Jun 2026 19:31:33 -0400 Subject: [PATCH 5/5] refactor(magics): replace fragile psutil socket check with client close mock spy --- .../tests/system/test_bigquery.py | 82 +++++-------------- 1 file changed, 21 insertions(+), 61 deletions(-) diff --git a/packages/bigquery-magics/tests/system/test_bigquery.py b/packages/bigquery-magics/tests/system/test_bigquery.py index 441bc60924b2..0e97234eed1c 100644 --- a/packages/bigquery-magics/tests/system/test_bigquery.py +++ b/packages/bigquery-magics/tests/system/test_bigquery.py @@ -14,77 +14,42 @@ """System tests for Jupyter/IPython connector.""" -import contextlib -import gc import re -import time +from unittest import mock +import google.cloud.bigquery import pandas -import psutil from IPython.testing import globalipapp from IPython.utils import io -@contextlib.contextmanager -def patch_tracked_requests(): - """Context manager to patch google-auth requests and track/close their HTTP sessions. - - This prevents socket leaks in system tests that use Workload Identity or metadata server auth. - """ - import google.auth.transport.requests - - original_init = google.auth.transport.requests.Request.__init__ - tracked_requests = [] - - def patched_init(self, session=None): - original_init(self, session=session) - if session is None: - tracked_requests.append(self) - - google.auth.transport.requests.Request.__init__ = patched_init - try: - yield tracked_requests - finally: - google.auth.transport.requests.Request.__init__ = original_init - for req in tracked_requests: - if hasattr(req, "session") and req.session is not None: - req.session.close() - - def test_bigquery_magic(): globalipapp.start_ipython() ip = globalipapp.get_ipython() - current_process = psutil.Process() - # GC to ensure clean starting state - gc.collect() - conn_count_start = len(current_process.net_connections()) + ip.extension_manager.load_extension("bigquery_magics") + sql = """ + SELECT + CONCAT( + 'https://stackoverflow.com/questions/', + CAST(id as STRING)) as url, + view_count + FROM `bigquery-public-data.stackoverflow.posts_questions` + WHERE tags like '%google-bigquery%' + ORDER BY view_count DESC + LIMIT 10 + """ + original_close = google.cloud.bigquery.Client.close + with mock.patch.object( + google.cloud.bigquery.Client, "close", autospec=True + ) as mock_close: + mock_close.side_effect = original_close - with patch_tracked_requests(): - ip.extension_manager.load_extension("bigquery_magics") - sql = """ - SELECT - CONCAT( - 'https://stackoverflow.com/questions/', - CAST(id as STRING)) as url, - view_count - FROM `bigquery-public-data.stackoverflow.posts_questions` - WHERE tags like '%google-bigquery%' - ORDER BY view_count DESC - LIMIT 10 - """ with io.capture_output() as captured: result = ip.run_cell_magic("bigquery", "--use_rest_api", sql) - # Force garbage collection to sweep unreferenced socket objects - gc.collect() - - # Wait a bit for the asynchronous channel teardown to complete and the socket to be closed. - for _ in range(30): - conn_count_end = len(current_process.net_connections()) - if conn_count_end <= conn_count_start: - break - time.sleep(0.1) + # Verify that client close is explicitly called to release sockets. + assert mock_close.called lines = re.split("\n|\r", captured.stdout) # Removes blanks & terminal code (result of display clearing) @@ -94,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