Skip to content

Commit 409d91b

Browse files
committed
add unit test
1 parent a610a63 commit 409d91b

File tree

3 files changed

+34
-6
lines changed

3 files changed

+34
-6
lines changed

bigframes/_config/bigquery_options.py

Lines changed: 11 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -396,6 +396,17 @@ def requests_transport_adapters(
396396
adapters
397397
<https://requests.readthedocs.io/en/latest/user/advanced/#transport-adapters>`_.
398398
399+
**Examples:**
400+
401+
Increase the connection pool size using the requests `HTTPAdapter
402+
<https://requests.readthedocs.io/en/latest/api/#requests.adapters.HTTPAdapter>`_.
403+
404+
>>> import bigframes.pandas as bpd
405+
>>> bpd.options.bigquery.requests_transport_adapters = (
406+
... ("http://", requests.adapters.HTTPAdapter(pool_maxsize=100)),
407+
... ("https://", requests.adapters.HTTPAdapter(pool_maxsize=100)),
408+
... ) # doctest: +SKIP
409+
399410
Returns:
400411
Sequence[Tuple[str, requests.adapters.BaseAdapter]]:
401412
Prefixes and corresponding transport adapters to `mount

tests/unit/_config/test_bigquery_options.py

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -38,6 +38,7 @@
3838
("skip_bq_connection_check", False, True),
3939
("client_endpoints_override", {}, {"bqclient": "endpoint_address"}),
4040
("ordering_mode", "strict", "partial"),
41+
("requests_transport_adapters", object(), object()),
4142
],
4243
)
4344
def test_setter_raises_if_session_started(attribute, original_value, new_value):

tests/unit/session/test_clients.py

Lines changed: 22 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -13,25 +13,22 @@
1313
# limitations under the License.
1414

1515
import os
16-
from typing import Optional
16+
from typing import cast, Optional
1717
import unittest.mock as mock
1818

19-
import google.api_core.client_info
20-
import google.api_core.client_options
21-
import google.api_core.exceptions
22-
import google.api_core.gapic_v1.client_info
2319
import google.auth.credentials
2420
import google.cloud.bigquery
2521
import google.cloud.bigquery_connection_v1
2622
import google.cloud.bigquery_storage_v1
2723
import google.cloud.functions_v2
2824
import google.cloud.resourcemanager_v3
25+
import requests.adapters
2926

3027
import bigframes.session.clients as clients
3128
import bigframes.version
3229

3330

34-
def create_clients_provider(application_name: Optional[str] = None):
31+
def create_clients_provider(application_name: Optional[str] = None, **kwargs):
3532
credentials = mock.create_autospec(google.auth.credentials.Credentials)
3633
return clients.ClientsProvider(
3734
project="test-project",
@@ -40,6 +37,7 @@ def create_clients_provider(application_name: Optional[str] = None):
4037
credentials=credentials,
4138
application_name=application_name,
4239
bq_kms_key_name="projects/my-project/locations/us/keyRings/myKeyRing/cryptoKeys/myKey",
40+
**kwargs,
4341
)
4442

4543

@@ -127,6 +125,24 @@ def assert_clients_wo_user_agent(
127125
)
128126

129127

128+
def test_requests_transport_adapters_pool_maxsize(monkeypatch):
129+
monkeypatch_client_constructors(monkeypatch)
130+
requests_transport_adapters = (
131+
("http://", requests.adapters.HTTPAdapter(pool_maxsize=123)),
132+
("https://", requests.adapters.HTTPAdapter(pool_maxsize=123)),
133+
) # doctest: +SKIP
134+
provider = create_clients_provider(
135+
requests_transport_adapters=requests_transport_adapters
136+
)
137+
138+
_, kwargs = cast(mock.Mock, provider.bqclient).call_args
139+
requests_session = kwargs.get("_http")
140+
adapter: requests.adapters.HTTPAdapter = requests_session.get_adapter(
141+
"https://bigquery.googleapis.com/"
142+
)
143+
assert adapter._pool_maxsize == 123 # type: ignore
144+
145+
130146
def test_user_agent_default(monkeypatch):
131147
monkeypatch_client_constructors(monkeypatch)
132148
provider = create_clients_provider(application_name=None)

0 commit comments

Comments
 (0)