From 83d9e3f25f9edf5817cf72b83318132e66522f4d Mon Sep 17 00:00:00 2001 From: Carmen Chow Date: Thu, 11 Jul 2024 15:55:02 -0700 Subject: [PATCH 1/3] add additional get_aws_credentials tests --- tests/cli/test_cmds_spark_run.py | 72 ++++++++++++++++++++++++++++++++ 1 file changed, 72 insertions(+) diff --git a/tests/cli/test_cmds_spark_run.py b/tests/cli/test_cmds_spark_run.py index 42a5445a68..be57b833a0 100644 --- a/tests/cli/test_cmds_spark_run.py +++ b/tests/cli/test_cmds_spark_run.py @@ -1450,3 +1450,75 @@ def test_get_aws_credentials(): RoleSessionName=mock.ANY, WebIdentityToken="token-content", ) + + +@mock.patch("service_configuration_lib.spark_config.use_aws_profile", autospec=False) +@mock.patch("service_configuration_lib.spark_config.Session", autospec=True) +def test_get_aws_credentials_session(mock_boto3_session, mock_use_aws_profile): + # prioritize session over `profile_name` if both are provided + from service_configuration_lib.spark_config import get_aws_credentials + + session = mock_boto3_session() + + get_aws_credentials( + service="some-service", + session=session, + profile_name="some-profile", + ) + + mock_use_aws_profile.assert_called_once_with(session=session) + session.assert_not_called() + + +@mock.patch("service_configuration_lib.spark_config.use_aws_profile", autospec=False) +@mock.patch("service_configuration_lib.spark_config.Session", autospec=True) +def test_get_aws_credentials_profile(mock_boto3_session, mock_use_aws_profile): + # prioritize `profile_name` over `service` if both are provided + from service_configuration_lib.spark_config import get_aws_credentials + + profile_name = "some-profile" + + get_aws_credentials(service="some-service", profile_name=profile_name) + + mock_use_aws_profile.assert_called_once_with(profile_name=profile_name) + + +@mock.patch("service_configuration_lib.spark_config.use_aws_profile", autospec=False) +@mock.patch("os.path.exists", autospec=False) +@mock.patch( + "service_configuration_lib.spark_config._load_aws_credentials_from_yaml", + autospec=True, +) +def test_get_aws_credentials_boto_cfg( + mock_load_aws_credentials_from_yaml, mock_os_path_exists, mock_use_aws_profile +): + # use `service` if profile_name is not provided + from service_configuration_lib.spark_config import ( + get_aws_credentials, + AWS_CREDENTIALS_DIR, + ) + + service_name = "some-service" + + get_aws_credentials( + service=service_name, + ) + + credentials_path = f"{AWS_CREDENTIALS_DIR}{service_name}.yaml" + mock_os_path_exists.return_value = True + + mock_load_aws_credentials_from_yaml.assert_called_once_with(credentials_path) + mock_use_aws_profile.assert_not_called() + + +@mock.patch("service_configuration_lib.spark_config.use_aws_profile", autospec=False) +@mock.patch("service_configuration_lib.spark_config.Session", autospec=True) +def test_get_aws_credentials_default_profile(mock_boto3_session, mock_use_aws_profile): + # use `default` profile if no valid options are provided + from service_configuration_lib.spark_config import get_aws_credentials + + get_aws_credentials( + service="spark", + ) + + mock_use_aws_profile.assert_called_once() From 3dbed5325cebf67f225c32ff4b0e18b86c5f986f Mon Sep 17 00:00:00 2001 From: Carmen Chow Date: Wed, 17 Jul 2024 09:26:25 -0700 Subject: [PATCH 2/3] move import statements to the top --- tests/cli/test_cmds_spark_run.py | 17 +++-------------- 1 file changed, 3 insertions(+), 14 deletions(-) diff --git a/tests/cli/test_cmds_spark_run.py b/tests/cli/test_cmds_spark_run.py index be57b833a0..5f82a6dcfd 100644 --- a/tests/cli/test_cmds_spark_run.py +++ b/tests/cli/test_cmds_spark_run.py @@ -12,9 +12,12 @@ # See the License for the specific language governing permissions and # limitations under the License. import argparse +import os import mock import pytest +from service_configuration_lib.spark_config import AWS_CREDENTIALS_DIR +from service_configuration_lib.spark_config import get_aws_credentials from paasta_tools import spark_tools from paasta_tools import utils @@ -1422,9 +1425,6 @@ def test_build_and_push_docker_image_unexpected_output_format( def test_get_aws_credentials(): - import os - from service_configuration_lib.spark_config import get_aws_credentials - with mock.patch.dict( os.environ, { @@ -1456,8 +1456,6 @@ def test_get_aws_credentials(): @mock.patch("service_configuration_lib.spark_config.Session", autospec=True) def test_get_aws_credentials_session(mock_boto3_session, mock_use_aws_profile): # prioritize session over `profile_name` if both are provided - from service_configuration_lib.spark_config import get_aws_credentials - session = mock_boto3_session() get_aws_credentials( @@ -1474,8 +1472,6 @@ def test_get_aws_credentials_session(mock_boto3_session, mock_use_aws_profile): @mock.patch("service_configuration_lib.spark_config.Session", autospec=True) def test_get_aws_credentials_profile(mock_boto3_session, mock_use_aws_profile): # prioritize `profile_name` over `service` if both are provided - from service_configuration_lib.spark_config import get_aws_credentials - profile_name = "some-profile" get_aws_credentials(service="some-service", profile_name=profile_name) @@ -1493,11 +1489,6 @@ def test_get_aws_credentials_boto_cfg( mock_load_aws_credentials_from_yaml, mock_os_path_exists, mock_use_aws_profile ): # use `service` if profile_name is not provided - from service_configuration_lib.spark_config import ( - get_aws_credentials, - AWS_CREDENTIALS_DIR, - ) - service_name = "some-service" get_aws_credentials( @@ -1515,8 +1506,6 @@ def test_get_aws_credentials_boto_cfg( @mock.patch("service_configuration_lib.spark_config.Session", autospec=True) def test_get_aws_credentials_default_profile(mock_boto3_session, mock_use_aws_profile): # use `default` profile if no valid options are provided - from service_configuration_lib.spark_config import get_aws_credentials - get_aws_credentials( service="spark", ) From 763ac6e76f0fb7179fd7f813941554f5e8b34be0 Mon Sep 17 00:00:00 2001 From: Carmen Date: Thu, 18 Jul 2024 11:06:31 -0700 Subject: [PATCH 3/3] Update tests/cli/test_cmds_spark_run.py Co-authored-by: Sameer Sharma <30409342+CaptainSame@users.noreply.github.com> --- tests/cli/test_cmds_spark_run.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tests/cli/test_cmds_spark_run.py b/tests/cli/test_cmds_spark_run.py index 5f82a6dcfd..97f81dd432 100644 --- a/tests/cli/test_cmds_spark_run.py +++ b/tests/cli/test_cmds_spark_run.py @@ -1510,4 +1510,4 @@ def test_get_aws_credentials_default_profile(mock_boto3_session, mock_use_aws_pr service="spark", ) - mock_use_aws_profile.assert_called_once() + mock_use_aws_profile.assert_called_once_with()