"""Test assume_role."""
from unittest.mock import patch

import pytest

from feed_ingestion.util.aws import assume_role


_FAKE_CREDENTIALS = {
    'Credentials': {
        'AccessKeyId': 'FAKE_KEY',
        'SecretAccessKey': 'FAKE_SECRET',
        'SessionToken': 'FAKE_TOKEN',
    }
}


@pytest.fixture
def mock_boto3():
    """Mock boto3."""
    boto3_path = 'feed_ingestion.util.aws.assume_role.boto3'
    with patch(boto3_path) as boto3:
        sts_client = boto3.session.Session.return_value.client.return_value
        sts_client.assume_role.return_value = _FAKE_CREDENTIALS
        yield boto3


@pytest.fixture
def sts_client(mock_boto3):
    """Return the mocked STS client."""
    return mock_boto3.session.Session.return_value.client.return_value


def test_assumed_session(mock_boto3, sts_client):
    """Test assumed_session assumes role and returns Session with creds."""
    result = assume_role.assumed_session('arn:aws:iam::123:role/test')

    sts_client.assume_role.assert_called_once_with(
        RoleArn='arn:aws:iam::123:role/test',
        RoleSessionName='feed_ingestion')
    mock_boto3.Session.assert_called_once_with(
        aws_access_key_id='FAKE_KEY',
        aws_secret_access_key='FAKE_SECRET',
        aws_session_token='FAKE_TOKEN',
    )
    assert result == mock_boto3.Session.return_value


def test_assumed_session_custom_session_name(mock_boto3, sts_client):
    """Test assumed_session passes custom session_name to assume_role."""
    assume_role.assumed_session('arn:aws:iam::123:role/test', 'my_session')

    sts_client.assume_role.assert_called_once_with(
        RoleArn='arn:aws:iam::123:role/test',
        RoleSessionName='my_session')


def test_get_s3_client_assume_role(mock_boto3, sts_client):
    """Test get_s3_client_assume_role assumes role and returns S3 client."""
    result = assume_role.get_s3_client_assume_role('test_role')

    sts_client.assume_role.assert_called_once_with(
        RoleArn='test_role', RoleSessionName='feed_ingestion')
    mock_boto3.Session.return_value.client.assert_called_once_with('s3')
    assert result == mock_boto3.Session.return_value.client.return_value