Make sure that an arn including 'role' won't pass (#65) (#570)

Author: laurenyu

CHANGELOG.rst

@@ -2,10 +2,15 @@
 CHANGELOG
 =========
 
+1.16.4.dev
+==========
+
+* bug-fix: Session: don't allow get_execution_role() to return an ARN that's not a role but has "role" in the name
+
 1.16.3
 ======
 
-* bug-fix: Local Mode: Allow support for SSH in local mode 
+* bug-fix: Local Mode: Allow support for SSH in local mode
 * bug-fix: Append retry id to default Airflow job name to avoid name collisions in retry
 * bug-fix: Local Mode: No longer requires s3 permissions to run local entry point file
 * feature: Estimators: add support for PyTorch 1.0.0

src/sagemaker/session.py

@@ -1212,7 +1212,7 @@ def get_execution_role(sagemaker_session=None):
         sagemaker_session = Session()
     arn = sagemaker_session.get_caller_identity_arn()
 
-    if 'role' in arn:
+    if ':role/' in arn:
         return arn
     message = 'The current AWS identity is not a role: {}, therefore it cannot be used as a SageMaker execution role'
     raise ValueError(message.format(arn))

tests/unit/test_session.py

@@ -65,6 +65,15 @@ def test_get_execution_role_throws_exception_if_arn_is_not_role():
     assert 'ValueError: The current AWS identity is not a role' in str(error)
 
 
+def test_get_execution_role_throws_exception_if_arn_is_not_role_with_role_in_name():
+    session = Mock()
+    session.get_caller_identity_arn.return_value = 'arn:aws:iam::369233609183:user/marcos-role'
+
+    with pytest.raises(ValueError) as error:
+        get_execution_role(session)
+    assert 'ValueError: The current AWS identity is not a role' in str(error)
+
+
 def test_get_caller_identity_arn_from_an_user(boto_session):
     sess = Session(boto_session)
     arn = 'arn:aws:iam::369233609183:user/mia'