Add tests

Author: trungleduc

src/sagemaker/session.py

@@ -6862,8 +6862,8 @@ def get_execution_role(sagemaker_session=None, use_default=False):
     Throws an exception if role doesn't exist.
 
     Args:
-        sagemaker_session(Session): Current sagemaker session.
-        use_default(bool): Use a default role if `get_caller_identity_arn does not
+        sagemaker_session (Session): Current sagemaker session.
+        use_default (bool): Use a default role if ``get_caller_identity_arn`` does not
             return a correct role. This default role will be created if needed.
             Defaults to ``False``.
 

tests/unit/test_session.py

@@ -15,6 +15,7 @@
 import copy
 import datetime
 import io
+import json
 import logging
 import os
 
@@ -532,6 +533,67 @@ def test_get_execution_role_throws_exception_if_arn_is_not_role_with_role_in_nam
     assert "The current AWS identity is not a role" in str(error.value)
 
 
+def test_get_execution_role_get_default_role(caplog):
+    session = Mock()
+    session.get_caller_identity_arn.return_value = "arn:aws:iam::369233609183:user/marcos"
+
+    iam_client = Mock()
+    iam_client.get_role.return_value = {"Role": {"Arn": "foo-role"}}
+    boto_session = Mock()
+    boto_session.client.return_value = iam_client
+
+    session.boto_session = boto_session
+    actual = get_execution_role(session, use_default=True)
+
+    iam_client.get_role.assert_called_with(RoleName="AmazonSageMaker-DefaultRole")
+    iam_client.attach_role_policy.assert_called_with(
+        PolicyArn="arn:aws:iam::aws:policy/AmazonSageMakerFullAccess",
+        RoleName="AmazonSageMaker-DefaultRole",
+    )
+    assert "Using default role: AmazonSageMaker-DefaultRole" in caplog.text
+    assert actual == "foo-role"
+
+
+def test_get_execution_role_create_default_role(caplog):
+    session = Mock()
+    session.get_caller_identity_arn.return_value = "arn:aws:iam::369233609183:user/marcos"
+    permissions_policy = json.dumps(
+        {
+            "Version": "2012-10-17",
+            "Statement": [
+                {
+                    "Effect": "Allow",
+                    "Principal": {"Service": ["sagemaker.amazonaws.com"]},
+                    "Action": "sts:AssumeRole",
+                }
+            ],
+        }
+    )
+    iam_client = Mock()
+    iam_client.exceptions.NoSuchEntityException = Exception
+    iam_client.get_role = Mock(side_effect=[Exception(), {"Role": {"Arn": "foo-role"}}])
+
+    boto_session = Mock()
+    boto_session.client.return_value = iam_client
+
+    session.boto_session = boto_session
+
+    actual = get_execution_role(session, use_default=True)
+
+    iam_client.create_role.assert_called_with(
+        RoleName="AmazonSageMaker-DefaultRole", AssumeRolePolicyDocument=str(permissions_policy)
+    )
+
+    iam_client.attach_role_policy.assert_called_with(
+        PolicyArn="arn:aws:iam::aws:policy/AmazonSageMakerFullAccess",
+        RoleName="AmazonSageMaker-DefaultRole",
+    )
+
+    assert "Created new sagemaker execution role: AmazonSageMaker-DefaultRole" in caplog.text
+
+    assert actual == "foo-role"
+
+
 @patch(
     "six.moves.builtins.open",
     mock_open(read_data='{"ResourceName": "SageMakerInstance"}'),