diff --git a/packages/google-cloud-bigquery/tests/system/test_client.py b/packages/google-cloud-bigquery/tests/system/test_client.py index d5ec07b5a557..dc3f645dbce2 100644 --- a/packages/google-cloud-bigquery/tests/system/test_client.py +++ b/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 diff --git a/packages/google-cloud-bigquery/tests/system/test_magics.py b/packages/google-cloud-bigquery/tests/system/test_magics.py index 31fd4543eed5..85ca0a243419 100644 --- a/packages/google-cloud-bigquery/tests/system/test_magics.py +++ b/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)