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

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,15 @@ 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)