Mock assumed role for global default session.

[bonext]

Hello,

I'm trying to test some SageMaker code that makes use of sagemaker.get_execution_role().

When called without parameters, that ends up using the default boto session to run sts get_caller_identity().

It is my understanding (from the following SDK code) that in any reasonable sagemaker usage this should return an assumed role arn, but when using moto I always get the default moto user arn.

Is it possible to set up moto and/or my test mocks to get an assumed role arn as a caller identity from the default session?

[bonext]

Answering my own question, here's how.

import pytest
from moto import mock_aws
import json
 

@pytest.fixture()
def aws_credentials(monkeypatch):
    """Mocked AWS Credentials for moto."""
    monkeypatch.setenv("AWS_DEFAULT_REGION", "eu-west-2")
    monkeypatch.setenv("AWS_ACCESS_KEY_ID", "testing")
    monkeypatch.setenv("AWS_SECRET_ACCESS_KEY", "testing")
    monkeypatch.setenv("AWS_SECURITY_TOKEN", "testing")
    monkeypatch.setenv("AWS_SESSION_TOKEN", "testing")
 
 
@pytest.fixture()
def mocked_aws(aws_credentials):
    """
    Mock all AWS interactions
    Requires you to create your own boto3 clients
    """
    with mock_aws():
        yield


@pytest.fixture()
def mock_role_name():
    return "test_sagemaker_role"


@pytest.fixture()
def mock_sagemaker_role(mocked_aws, mock_role_name):
    import boto3
    iam = boto3.client("iam")
    resp = iam.create_role(
        RoleName=mock_role_name,
        AssumeRolePolicyDocument=json.dumps(
            {
                "Version": "2012-10-17",
                "Statement": [
                    {
                        "Effect": "Allow",
                        "Principal": {"Service": "sagemaker.amazonaws.com"},
                        "Action": "sts:AssumeRole",
                    }
                ],
            }
        ),
    )
    return resp


@pytest.fixture()
def mock_sagemaker_role_arn(mock_sagemaker_role):
    return mock_sagemaker_role["Role"]["Arn"]


@pytest.fixture()
def mock_assumed_sagemaker_role_arn(mock_sagemaker_role_arn):
    """emulate expected arn for assumed role based on sagemaker SDK code"""
    assumed_role_arn = mock_sagemaker_role_arn.replace(":role/", ":assumed-role/")
    return assumed_role_arn


@pytest.fixture()
def mock_assumed_role_session(mock_assumed_sagemaker_role_arn):
    import boto3

    sts = boto3.client("sts")
    response = sts.assume_role(
        RoleArn=mock_assumed_sagemaker_role_arn,
        RoleSessionName="test-session-name",
        ExternalId="test-external-id",
    )

    boto3.setup_default_session(
        aws_access_key_id=response["Credentials"]["AccessKeyId"],
        aws_secret_access_key=response["Credentials"]["SecretAccessKey"],
        aws_session_token=response["Credentials"]["SessionToken"],
    )

     
def test_sagemaker_get_execution_role(mock_assumed_role_session, mock_sagemaker_role_arn):
    import sagemaker
    assert sagemaker.get_execution_role() == mock_sagemaker_role_arn

[bblommers]

Thank you for posting the answer @bonext!