Skip to content

Commit fb580fe

Browse files
kyleknapkyleknap
authored
Introduce connect client tests (#507)
1 parent 91e0ff3 commit fb580fe

File tree

5 files changed

+204
-84
lines changed

5 files changed

+204
-84
lines changed

‎CHANGELOG.md‎

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -4,6 +4,9 @@ Unreleased
44
----------
55

66
- Added "adlfs" to library's default user agent
7+
- Fix issue where ``AzureBlobFile`` did not respect ``location_mode`` parameter
8+
from parent ``AzureBlobFileSystem`` when using SAS credentials and connecting to
9+
new SDK clients.
710

811

912
2024.12.0

‎adlfs/spec.py‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -2093,6 +2093,7 @@ def connect_client(self):
20932093
elif self.fs.sas_token is not None:
20942094
self.container_client = _create_aio_blob_service_client(
20952095
account_url=self.fs.account_url + self.fs.sas_token,
2096+
location_mode=self.fs.location_mode,
20962097
).get_container_client(self.container_name)
20972098
else:
20982099
self.container_client = _create_aio_blob_service_client(

‎adlfs/tests/constants.py‎

Lines changed: 3 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1,6 +1,8 @@
1-
URL = "http://127.0.0.1:10000"
1+
HOST = "127.0.0.1:10000"
2+
URL = f"http://{HOST}"
23
ACCOUNT_NAME = "devstoreaccount1"
34
KEY = "Eby8vdM02xNOcqFlqUwJPLlmEtlCDXJ1OUzFT50uSRZ6IFsuFq2UVErCz4I6tq/K1SZFPTOtr/KBHBeksoGMGw==" # NOQA
45
CONN_STR = f"DefaultEndpointsProtocol=http;AccountName={ACCOUNT_NAME};AccountKey={KEY};BlobEndpoint={URL}/{ACCOUNT_NAME};" # NOQA
6+
SAS_TOKEN = "not-a-real-sas-token"
57
DEFAULT_VERSION_ID = "1970-01-01T00:00:00.0000000Z"
68
LATEST_VERSION_ID = "2022-01-01T00:00:00.0000000Z"

‎adlfs/tests/test_connect_client.py‎

Lines changed: 197 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,197 @@
1+
from unittest import mock
2+
3+
import azure.storage.blob
4+
import pytest
5+
from azure.storage.blob.aio import BlobServiceClient as AIOBlobServiceClient
6+
7+
from adlfs import AzureBlobFile, AzureBlobFileSystem
8+
from adlfs.tests.constants import ACCOUNT_NAME, CONN_STR, HOST, KEY, SAS_TOKEN
9+
from adlfs.utils import __version__ as __version__
10+
11+
12+
@pytest.fixture()
13+
def mock_from_connection_string(mocker):
14+
return mocker.patch.object(
15+
AIOBlobServiceClient,
16+
"from_connection_string",
17+
autospec=True,
18+
side_effect=AIOBlobServiceClient.from_connection_string,
19+
)
20+
21+
22+
@pytest.fixture()
23+
def mock_service_client_init(mocker):
24+
return mocker.patch.object(
25+
AIOBlobServiceClient,
26+
"__init__",
27+
autospec=True,
28+
side_effect=AIOBlobServiceClient.__init__,
29+
)
30+
31+
32+
def get_expected_client_init_call(
33+
account_url,
34+
credential=None,
35+
location_mode="primary",
36+
):
37+
call_kwargs = {
38+
"account_url": account_url,
39+
"user_agent": f"adlfs/{__version__}",
40+
}
41+
if credential is not None:
42+
call_kwargs["credential"] = credential
43+
if location_mode is not None:
44+
call_kwargs["_location_mode"] = location_mode
45+
return mock.call(mock.ANY, **call_kwargs)
46+
47+
48+
def get_expected_client_from_connection_string_call(
49+
conn_str,
50+
):
51+
return mock.call(conn_str=conn_str, user_agent=f"adlfs/{__version__}")
52+
53+
54+
def assert_client_create_calls(
55+
mock_client_create_method,
56+
expected_create_call,
57+
expected_call_count=1,
58+
):
59+
expected_call_args_list = [expected_create_call for _ in range(expected_call_count)]
60+
assert mock_client_create_method.call_args_list == expected_call_args_list
61+
62+
63+
def ensure_no_api_calls_on_close(file_obj):
64+
# Marks the file-like object as closed to prevent any API calls during an invocation of
65+
# close(), which can occur during garbage collection or direct invocation.
66+
#
67+
# This is important for test cases where we do not want to make an API request whether:
68+
#
69+
# * The test would hang because Azurite is not configured to use SSL and adlfs always sets SSL
70+
# for SDK clients created via their initializer.
71+
#
72+
# * The filesystem is configured to use the secondary location which can by-pass the location
73+
# that Azurite is running.
74+
file_obj.closed = True
75+
76+
77+
@pytest.mark.parametrize(
78+
"fs_kwargs,expected_client_init_call",
79+
[
80+
(
81+
{"account_name": ACCOUNT_NAME, "account_key": KEY},
82+
get_expected_client_init_call(
83+
account_url=f"https://{ACCOUNT_NAME}.blob.core.windows.net",
84+
credential=KEY,
85+
),
86+
),
87+
(
88+
{"account_name": ACCOUNT_NAME, "credential": SAS_TOKEN},
89+
get_expected_client_init_call(
90+
account_url=f"https://{ACCOUNT_NAME}.blob.core.windows.net",
91+
credential=SAS_TOKEN,
92+
),
93+
),
94+
(
95+
{"account_name": ACCOUNT_NAME, "sas_token": SAS_TOKEN},
96+
get_expected_client_init_call(
97+
account_url=f"https://{ACCOUNT_NAME}.blob.core.windows.net?{SAS_TOKEN}",
98+
),
99+
),
100+
# Anonymous connection
101+
(
102+
{"account_name": ACCOUNT_NAME},
103+
get_expected_client_init_call(
104+
account_url=f"https://{ACCOUNT_NAME}.blob.core.windows.net",
105+
location_mode=None,
106+
),
107+
),
108+
# Override host
109+
(
110+
{
111+
"account_name": ACCOUNT_NAME,
112+
"account_host": HOST,
113+
"sas_token": SAS_TOKEN,
114+
},
115+
get_expected_client_init_call(
116+
account_url=f"https://{HOST}?{SAS_TOKEN}",
117+
),
118+
),
119+
# Override location mode
120+
(
121+
{
122+
"account_name": ACCOUNT_NAME,
123+
"account_key": KEY,
124+
"location_mode": "secondary",
125+
},
126+
get_expected_client_init_call(
127+
account_url=f"https://{ACCOUNT_NAME}.blob.core.windows.net",
128+
credential=KEY,
129+
location_mode="secondary",
130+
),
131+
),
132+
(
133+
{
134+
"account_name": ACCOUNT_NAME,
135+
"credential": SAS_TOKEN,
136+
"location_mode": "secondary",
137+
},
138+
get_expected_client_init_call(
139+
account_url=f"https://{ACCOUNT_NAME}.blob.core.windows.net",
140+
credential=SAS_TOKEN,
141+
location_mode="secondary",
142+
),
143+
),
144+
(
145+
{
146+
"account_name": ACCOUNT_NAME,
147+
"sas_token": SAS_TOKEN,
148+
"location_mode": "secondary",
149+
},
150+
get_expected_client_init_call(
151+
account_url=f"https://{ACCOUNT_NAME}.blob.core.windows.net?{SAS_TOKEN}",
152+
location_mode="secondary",
153+
),
154+
),
155+
],
156+
)
157+
def test_connect_initializer(
158+
storage: azure.storage.blob.BlobServiceClient,
159+
mock_service_client_init,
160+
fs_kwargs,
161+
expected_client_init_call,
162+
):
163+
fs = AzureBlobFileSystem(skip_instance_cache=True, **fs_kwargs)
164+
assert_client_create_calls(mock_service_client_init, expected_client_init_call)
165+
166+
f = AzureBlobFile(fs, "data/root/a/file.txt", mode="wb")
167+
f.connect_client()
168+
ensure_no_api_calls_on_close(f)
169+
assert_client_create_calls(
170+
mock_service_client_init,
171+
expected_client_init_call,
172+
expected_call_count=2,
173+
)
174+
175+
176+
def test_connect_connection_str(
177+
storage: azure.storage.blob.BlobServiceClient, mock_from_connection_string
178+
):
179+
fs = AzureBlobFileSystem(
180+
account_name=storage.account_name,
181+
connection_string=CONN_STR,
182+
skip_instance_cache=True,
183+
)
184+
expected_from_connection_str_call = get_expected_client_from_connection_string_call(
185+
conn_str=CONN_STR,
186+
)
187+
assert_client_create_calls(
188+
mock_from_connection_string, expected_from_connection_str_call
189+
)
190+
191+
f = AzureBlobFile(fs, "data/root/a/file.txt", mode="rb")
192+
f.connect_client()
193+
assert_client_create_calls(
194+
mock_from_connection_string,
195+
expected_from_connection_str_call,
196+
expected_call_count=2,
197+
)

‎adlfs/tests/test_user_agent.py‎

Lines changed: 0 additions & 83 deletions
This file was deleted.

0 commit comments

Comments
 (0)