tests(bigquery): implement robust wait loop for socket leak tests (#17688)

## Intent
This PR addresses flakiness in the `google-cloud-bigquery` system tests
that verify system resource cleanup (specifically open sockets).

## The Problem
Some system tests in `google-cloud-bigquery` were performing
instantaneous assertions on the number of open sockets immediately after
client closure or garbage collection. These assertions are fragile
because underlying transport layers (like gRPC) and OS socket
reclamation can occur asynchronously in background threads. This leads
to intermittent failures (flakiness) especially under varied
environmental loads or parallel execution.

## Solution
Refactored the fragile instantaneous checks in `test_client.py` and
`test_magics.py` to use a **Robust Wait-Loop** pattern.
- The tests now wait up to 3 seconds, checking periodically, to allow
asynchronous cleanup to complete before asserting.
- This pattern is already proven effective in other packages like
`google-cloud-pubsub`.
- Added explicit `gc.collect()` to `test_magics.py` to align with best
practices for resource verification tests.

> [!note]
> Also includes some linting updates. Sorry.

Author: chalmerlowe

packages/google-cloud-bigquery/tests/system/test_client.py

@@ -22,38 +22,36 @@
 import operator
 import os
 import pathlib
+import random
+import string
 import time
 import unittest
 import uuid
-import random
-import string
 from typing import Optional
 
-from google.api_core.exceptions import PreconditionFailed
-from google.api_core.exceptions import BadRequest
-from google.api_core.exceptions import ClientError
-from google.api_core.exceptions import Conflict
-from google.api_core.exceptions import GoogleAPICallError
-from google.api_core.exceptions import NotFound
-from google.api_core.exceptions import InternalServerError
-from google.api_core.exceptions import ServiceUnavailable
-from google.api_core.exceptions import TooManyRequests
-from google.cloud import bigquery
-from google.cloud.bigquery.dataset import Dataset
-from google.cloud.bigquery.dataset import DatasetReference
-from google.cloud.bigquery.table import Table
+import psutil
+import pytest
+from google.api_core.exceptions import (
+    BadRequest,
+    ClientError,
+    Conflict,
+    GoogleAPICallError,
+    InternalServerError,
+    NotFound,
+    PreconditionFailed,
+    ServiceUnavailable,
+    TooManyRequests,
+)
+from google.cloud import bigquery, storage
 from google.cloud._helpers import UTC
 from google.cloud.bigquery import dbapi, enums
-from google.cloud import storage
-from google.cloud.datacatalog_v1 import types as datacatalog_types
+from google.cloud.bigquery.dataset import Dataset, DatasetReference
+from google.cloud.bigquery.table import Table
 from google.cloud.datacatalog_v1 import PolicyTagManagerClient
-from google.cloud.resourcemanager_v3 import types as resourcemanager_types
+from google.cloud.datacatalog_v1 import types as datacatalog_types
 from google.cloud.resourcemanager_v3 import TagKeysClient, TagValuesClient
-import psutil
-import pytest
-from test_utils.retry import RetryErrors
-from test_utils.retry import RetryInstanceState
-from test_utils.retry import RetryResult
+from google.cloud.resourcemanager_v3 import types as resourcemanager_types
+from test_utils.retry import RetryErrors, RetryInstanceState, RetryResult
 from test_utils.system import unique_resource_id
 
 from . import helpers
@@ -254,8 +252,13 @@ def test_close_releases_open_sockets(self):
         import gc
 
         gc.collect()
-        conn_end = current_process.net_connections()
-        conn_count_end = len(conn_end)
+        for _ in range(30):  # Wait up to 3 seconds
+            conn_end = current_process.net_connections()
+            conn_count_end = len(conn_end)
+            if conn_count_end <= conn_count_start:
+                break
+            time.sleep(0.1)
+
         self.assertLessEqual(conn_count_end, conn_count_start)
 
     def test_create_dataset(self):
@@ -635,8 +638,7 @@ def test_create_table_with_default_value_expression(self):
         self.assertEqual("FOO", row_2.get("username"))
 
     def test_create_table_w_time_partitioning_w_clustering_fields(self):
-        from google.cloud.bigquery.table import TimePartitioning
-        from google.cloud.bigquery.table import TimePartitioningType
+        from google.cloud.bigquery.table import TimePartitioning, TimePartitioningType
 
         dataset = self.temp_dataset(_make_dataset_id("create_table_tp_cf"))
         table_id = "test_table"
@@ -964,12 +966,12 @@ def test_update_table_clustering_configuration(self):
         self.assertIsNone(table3.clustering_fields, None)
 
     def test_update_table_constraints(self):
-        from google.cloud.bigquery.table import TableConstraints
         from google.cloud.bigquery.table import (
-            PrimaryKey,
+            ColumnReference,
             ForeignKey,
+            PrimaryKey,
+            TableConstraints,
             TableReference,
-            ColumnReference,
         )
 
         dataset = self.temp_dataset(_make_dataset_id("update_table"))
@@ -1128,8 +1130,7 @@ def test_insert_rows_then_dump_table(self):
         self.assertEqual(sorted(row_tuples, key=by_age), sorted(ROWS, key=by_age))
 
     def test_load_table_from_local_avro_file_then_dump_table(self):
-        from google.cloud.bigquery.job import SourceFormat
-        from google.cloud.bigquery.job import WriteDisposition
+        from google.cloud.bigquery.job import SourceFormat, WriteDisposition
 
         TABLE_NAME = "test_table_avro"
         ROWS = [
@@ -1169,8 +1170,7 @@ def test_load_table_from_local_avro_file_then_dump_table(self):
 
     def test_load_table_from_local_parquet_file_decimal_types(self):
         from google.cloud.bigquery.enums import DecimalTargetType
-        from google.cloud.bigquery.job import SourceFormat
-        from google.cloud.bigquery.job import WriteDisposition
+        from google.cloud.bigquery.job import SourceFormat, WriteDisposition
 
         TABLE_NAME = "test_table_parquet"
 
@@ -1359,9 +1359,11 @@ def test_load_table_from_csv_w_picosecond_timestamp(self):
         self.assertEqual(table.num_rows, 3)
 
     def test_load_avro_from_uri_then_dump_table(self):
-        from google.cloud.bigquery.job import CreateDisposition
-        from google.cloud.bigquery.job import SourceFormat
-        from google.cloud.bigquery.job import WriteDisposition
+        from google.cloud.bigquery.job import (
+            CreateDisposition,
+            SourceFormat,
+            WriteDisposition,
+        )
 
         table_name = "test_table"
         rows = [
@@ -1399,9 +1401,11 @@ def test_load_avro_from_uri_then_dump_table(self):
         )
 
     def test_load_table_from_uri_then_dump_table(self):
-        from google.cloud.bigquery.job import CreateDisposition
-        from google.cloud.bigquery.job import SourceFormat
-        from google.cloud.bigquery.job import WriteDisposition
+        from google.cloud.bigquery.job import (
+            CreateDisposition,
+            SourceFormat,
+            WriteDisposition,
+        )
 
         TABLE_ID = "test_table"
         GS_URL = self._write_csv_to_storage(
@@ -2202,15 +2206,22 @@ def test_dbapi_connection_does_not_leak_sockets(self):
         import gc
 
         gc.collect()
-        conn_end = current_process.net_connections()
-        conn_count_end = len(conn_end)
+        for _ in range(30):  # Wait up to 3 seconds
+            conn_end = current_process.net_connections()
+            conn_count_end = len(conn_end)
+            if conn_count_end <= conn_count_start:
+                break
+            time.sleep(0.1)
+
         self.assertLessEqual(conn_count_end, conn_count_start)
 
     def _load_table_for_dml(self, rows, dataset_id, table_id):
         from google.cloud._testing import _NamedTemporaryFile
-        from google.cloud.bigquery.job import CreateDisposition
-        from google.cloud.bigquery.job import SourceFormat
-        from google.cloud.bigquery.job import WriteDisposition
+        from google.cloud.bigquery.job import (
+            CreateDisposition,
+            SourceFormat,
+            WriteDisposition,
+        )
 
         dataset = self.temp_dataset(dataset_id)
         greeting = bigquery.SchemaField("greeting", "STRING", mode="NULLABLE")
@@ -2732,8 +2743,7 @@ def test_nested_table_to_arrow(self):
         bigquery_storage = pytest.importorskip("google.cloud.bigquery_storage")
         pyarrow = pytest.importorskip("pyarrow")
         pyarrow.types = pytest.importorskip("pyarrow.types")
-        from google.cloud.bigquery.job import SourceFormat
-        from google.cloud.bigquery.job import WriteDisposition
+        from google.cloud.bigquery.job import SourceFormat, WriteDisposition
 
         SF = bigquery.SchemaField
         schema = [
@@ -2865,8 +2875,7 @@ def test_parameterized_types_round_trip(dataset_id: str):
 
 
 def test_table_snapshots(dataset_id: str):
-    from google.cloud.bigquery import CopyJobConfig
-    from google.cloud.bigquery import OperationType
+    from google.cloud.bigquery import CopyJobConfig, OperationType
 
     client = Config.CLIENT
 
@@ -2936,8 +2945,7 @@ def test_table_snapshots(dataset_id: str):
 
 
 def test_table_clones(dataset_id: str):
-    from google.cloud.bigquery import CopyJobConfig
-    from google.cloud.bigquery import OperationType
+    from google.cloud.bigquery import CopyJobConfig, OperationType
 
     client = Config.CLIENT
 

packages/google-cloud-bigquery/tests/system/test_magics.py

@@ -15,9 +15,10 @@
 """System tests for Jupyter/IPython connector."""
 
 import re
+import time
 
-import pytest
 import psutil
+import pytest
 
 from . import helpers
 
@@ -71,8 +72,16 @@ def test_bigquery_magic(ipython_interactive):
         with io.capture_output() as captured:
             result = ip.run_cell_magic("bigquery", "--use_rest_api", sql)
 
-    conn_end = current_process.net_connections()
-    conn_count_end = len(conn_end)
+    import gc
+
+    gc.collect()
+
+    for _ in range(30):  # Wait up to 3 seconds
+        conn_end = current_process.net_connections()
+        conn_count_end = len(conn_end)
+        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)