diff --git a/functions-python/helpers/tests/test_helpers.py b/functions-python/helpers/tests/test_helpers.py index 8f1dda004..a4a8e73d5 100644 --- a/functions-python/helpers/tests/test_helpers.py +++ b/functions-python/helpers/tests/test_helpers.py @@ -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" @@ -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( @@ -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" @@ -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, ) diff --git a/functions-python/helpers/utils.py b/functions-python/helpers/utils.py index 0bffa19a4..9a2eea72e 100644 --- a/functions-python/helpers/utils.py +++ b/functions-python/helpers/utils.py @@ -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(