Skip to content

Commit

Permalink
feat: fix download authentication feeds (#867)
Browse files Browse the repository at this point in the history
  • Loading branch information
qcdyx authored Feb 3, 2025
1 parent 5c62297 commit 4a0ba9f
Show file tree
Hide file tree
Showing 2 changed files with 24 additions and 10 deletions.
26 changes: 20 additions & 6 deletions functions-python/helpers/tests/test_helpers.py
Original file line number Diff line number Diff line change
Expand Up @@ -65,8 +65,14 @@ def test_download_and_get_hash(self):
if os.path.exists(file_path):
os.remove(file_path)

def test_download_and_get_hash_auth_type_1(self):
mock_binary_data = b"binary data for auth type 1"
def test_download_and_get_hash_auth_type_header(self):
"""
Test the download_and_get_hash function for authentication type 2 (headers).
This test verifies that the download_and_get_hash function correctly handles authentication type 2,
where the credentials are passed in the headers. It mocks the necessary components and checks that
the request is made with the appropriate headers.
"""
mock_binary_data = b"binary data for auth type 2"
expected_hash = hashlib.sha256(mock_binary_data).hexdigest()
file_path = "test_file.txt"
url = "https://test.com"
Expand All @@ -81,7 +87,7 @@ def test_download_and_get_hash_auth_type_1(self):
"urllib3.PoolManager.request", return_value=mock_response
) as mock_request:
result_hash = download_and_get_hash(
url, file_path, "sha256", 8192, 1, api_key_parameter_name, credentials
url, file_path, "sha256", 8192, 2, api_key_parameter_name, credentials
)

self.assertEqual(
Expand All @@ -104,8 +110,16 @@ def test_download_and_get_hash_auth_type_1(self):
if os.path.exists(file_path):
os.remove(file_path)

def test_download_and_get_hash_auth_type_2(self):
mock_binary_data = b"binary data for auth type 2"
def test_download_and_get_hash_auth_type_api_key(self):
"""
Test the download_and_get_hash function for authentication type 1 (API key).
This test verifies that the download_and_get_hash function correctly handles authentication type 1,
where the credentials are passed as a query parameter in the URL. It mocks the necessary components
and checks that the request is made with the appropriate URL containing the API key.
"""
mock_binary_data = b"binary data for auth type 1"
expected_hash = hashlib.sha256(mock_binary_data).hexdigest()
file_path = "test_file.txt"
base_url = "https://test.com"
Expand All @@ -126,7 +140,7 @@ def test_download_and_get_hash_auth_type_2(self):
file_path,
"sha256",
8192,
2,
1,
api_key_parameter_name,
credentials,
)
Expand Down
8 changes: 4 additions & 4 deletions functions-python/helpers/utils.py
Original file line number Diff line number Diff line change
Expand Up @@ -100,18 +100,18 @@ def download_and_get_hash(
ctx.load_default_certs()
ctx.options |= 0x4 # ssl.OP_LEGACY_SERVER_CONNECT

# authentication_type == 1 -> the credentials are passed in the header
# authentication_type == 1 -> the credentials are passed in the url
headers = {
"User-Agent": "Mozilla/5.0 (Linux; Android 6.0; Nexus 5 Build/MRA58N) "
"AppleWebKit/537.36 (KHTML, like Gecko) "
"Chrome/126.0.0.0 Mobile Safari/537.36"
}
if authentication_type == 1 and api_key_parameter_name and credentials:
headers[api_key_parameter_name] = credentials
url += f"?{api_key_parameter_name}={credentials}"

# authentication_type == 2 -> the credentials are passed in the url
# authentication_type == 2 -> the credentials are passed in the header
if authentication_type == 2 and api_key_parameter_name and credentials:
url += f"?{api_key_parameter_name}={credentials}"
headers[api_key_parameter_name] = credentials

with urllib3.PoolManager(ssl_context=ctx) as http:
with http.request(
Expand Down

0 comments on commit 4a0ba9f

Please sign in to comment.