def test_list_user_pagination(self, mock_cache, mock_boto3_client,
mock_creds):
# Arrange
UserFactory.create(username='******')
UserFactory.create(username='******')
SourceTableFactory(dataset=MasterDataSetFactory.create(
user_access_type='REQUIRES_AUTHENTICATION'))
mock_user_client = mock.Mock()
mock_user_client.list_users.side_effect = [
{
"UserList": [{
"Arn": "Arn",
"Email": "*****@*****.**",
"Role": "AUTHOR",
"UserName": "******",
}],
"NextToken":
"foo",
},
{
"UserList": [{
"Arn": "Arn2",
"Email": "*****@*****.**",
"Role": "AUTHOR",
"UserName": "******",
}]
},
]
mock_data_client = mock.Mock()
mock_sts_client = mock.Mock()
mock_boto3_client.side_effect = [
mock_user_client,
mock_data_client,
mock_sts_client,
]
mock_creds.return_value = [mock.Mock()]
# Act
sync_quicksight_permissions()
# Assert
assert mock_user_client.update_user.call_args_list == [
mock.call(
AwsAccountId=mock.ANY,
Namespace='default',
Role='AUTHOR',
CustomPermissionsName='author-custom-permissions',
UserName='******',
Email='*****@*****.**',
),
mock.call(
AwsAccountId=mock.ANY,
Namespace='default',
Role='AUTHOR',
CustomPermissionsName='author-custom-permissions',
UserName='******',
Email='*****@*****.**',
),
]
def test_missing_user_handled_gracefully(self, mock_cache,
mock_boto3_client, mock_creds):
# Arrange
user = UserFactory.create(username='******')
user2 = UserFactory.create(username='******')
SourceTableFactory(dataset=MasterDataSetFactory.create(
user_access_type='REQUIRES_AUTHENTICATION'))
mock_user_client = mock.Mock()
mock_user_client.describe_user.side_effect = [
botocore.exceptions.ClientError(
{
"Error": {
"Code": "ResourceNotFoundException",
"Message": "User not found",
}
},
'DescribeUser',
),
{
"User": {
"Arn": "Arn",
"Email": "*****@*****.**",
"Role": "AUTHOR"
}
},
]
mock_data_client = mock.Mock()
mock_sts_client = mock.Mock()
mock_boto3_client.side_effect = [
mock_user_client,
mock_data_client,
mock_sts_client,
]
# Act
sync_quicksight_permissions(user_sso_ids_to_update=[
str(user.profile.sso_id),
str(user2.profile.sso_id)
])
# Assert
assert mock_user_client.describe_user.call_args_list == [
mock.call(
AwsAccountId=mock.ANY,
Namespace='default',
UserName=f'quicksight_federation/{user.profile.sso_id}',
),
mock.call(
AwsAccountId=mock.ANY,
Namespace='default',
UserName=f'quicksight_federation/{user2.profile.sso_id}',
),
]
assert len(mock_data_client.create_data_source.call_args_list) == 1
assert len(mock_data_client.update_data_source.call_args_list) == 0
def handle(self, *args, **options):
sync_quicksight_permissions()
def test_create_new_data_source(self, mock_cache, mock_boto3_client,
mock_creds):
# Arrange
UserFactory.create(username='******')
SourceTableFactory(dataset=MasterDataSetFactory.create(
user_access_type='REQUIRES_AUTHENTICATION'))
mock_user_client = mock.Mock()
mock_user_client.list_users.return_value = {
"UserList": [{
"Arn": "Arn",
"Email": "*****@*****.**",
"Role": "AUTHOR",
"UserName": "******",
}]
}
mock_data_client = mock.Mock()
mock_sts_client = mock.Mock()
mock_boto3_client.side_effect = [
mock_user_client,
mock_data_client,
mock_sts_client,
]
mock_creds.return_value = [mock.Mock()]
# Act
sync_quicksight_permissions()
# Assert
assert mock_user_client.update_user.call_args_list == [
mock.call(
AwsAccountId=mock.ANY,
Namespace='default',
Role='AUTHOR',
CustomPermissionsName='author-custom-permissions',
UserName='******',
Email='*****@*****.**',
)
]
assert mock_data_client.create_data_source.call_args_list == [
mock.call(
AwsAccountId=mock.ANY,
DataSourceId=mock.ANY,
Name=mock.ANY,
DataSourceParameters={
'AuroraPostgreSqlParameters': {
'Host': mock.ANY,
'Port': mock.ANY,
'Database': mock.ANY,
}
},
Credentials={
'CredentialPair': {
'Username': mock.ANY,
'Password': mock.ANY
}
},
VpcConnectionProperties={'VpcConnectionArn': mock.ANY},
Type='AURORA_POSTGRESQL',
Permissions=[{
'Principal':
'Arn',
'Actions': [
'quicksight:DescribeDataSource',
'quicksight:DescribeDataSourcePermissions',
'quicksight:PassDataSource',
],
}],
)
]
assert mock_data_client.update_data_source.call_args_list == []
assert sorted(
mock_data_client.delete_data_source.call_args_list,
key=lambda x: x.kwargs['DataSourceId'],
) == [
mock.call(
AwsAccountId=mock.ANY,
DataSourceId='data-workspace-dev-my_database-88f3887d',
),
mock.call(
AwsAccountId=mock.ANY,
DataSourceId='data-workspace-dev-test_external_db2-88f3887d',
),
]
def test_poll_until_user_created(self, mock_cache, mock_boto3_client,
mock_creds):
# Arrange
user = UserFactory.create(username='******')
SourceTableFactory(dataset=MasterDataSetFactory.create(
user_access_type='REQUIRES_AUTHENTICATION'))
mock_user_client = mock.Mock()
mock_user_client.describe_user.side_effect = [
botocore.exceptions.ClientError(
{
"Error": {
"Code": "ResourceNotFoundException",
"Message": "User not found",
}
},
'DescribeUser',
),
] * 10 + [{
"User": {
"Arn": "Arn",
"Email": "*****@*****.**",
"Role": "AUTHOR",
"UserName": "******",
}
}]
mock_data_client = mock.Mock()
mock_sts_client = mock.Mock()
mock_boto3_client.side_effect = [
mock_user_client,
mock_data_client,
mock_sts_client,
]
# Act
with mock.patch('dataworkspace.apps.applications.utils.gevent.sleep'):
sync_quicksight_permissions(
user_sso_ids_to_update=[str(user.profile.sso_id)],
poll_for_user_creation=True,
)
# Assert
assert mock_user_client.update_user.call_args_list == [
mock.call(
AwsAccountId=mock.ANY,
Namespace='default',
Role='AUTHOR',
CustomPermissionsName='author-custom-permissions',
UserName='******',
Email='*****@*****.**',
)
]
assert (mock_user_client.describe_user.call_args_list == [
mock.call(
AwsAccountId=mock.ANY,
Namespace='default',
UserName=f'quicksight_federation/{user.profile.sso_id}',
),
] * 11)
assert len(mock_data_client.create_data_source.call_args_list) == 1
assert len(mock_data_client.update_data_source.call_args_list) == 0
def test_create_new_data_source(self, mock_cache, mock_boto3_client, mock_creds):
# Arrange
UserFactory.create(username="******")
SourceTableFactory(
dataset=MasterDataSetFactory.create(
user_access_type=UserAccessType.REQUIRES_AUTHENTICATION
)
)
mock_user_client = mock.Mock()
mock_user_client.list_users.return_value = {
"UserList": [
{
"Arn": "Arn",
"Email": "*****@*****.**",
"Role": "AUTHOR",
"UserName": "******",
}
]
}
mock_data_client = mock.Mock()
mock_sts_client = mock.Mock()
mock_boto3_client.side_effect = [
mock_user_client,
mock_data_client,
mock_sts_client,
]
mock_creds.return_value = [mock.Mock()]
# Act
sync_quicksight_permissions()
# Assert
assert mock_user_client.update_user.call_args_list == [
mock.call(
AwsAccountId=mock.ANY,
Namespace="default",
Role="AUTHOR",
CustomPermissionsName="author-custom-permissions",
UserName="******",
Email="*****@*****.**",
)
]
assert mock_data_client.create_data_source.call_args_list == [
mock.call(
AwsAccountId=mock.ANY,
DataSourceId=mock.ANY,
Name=mock.ANY,
DataSourceParameters={
"AuroraPostgreSqlParameters": {
"Host": mock.ANY,
"Port": mock.ANY,
"Database": mock.ANY,
}
},
Credentials={"CredentialPair": {"Username": mock.ANY, "Password": mock.ANY}},
VpcConnectionProperties={"VpcConnectionArn": mock.ANY},
Type="AURORA_POSTGRESQL",
Permissions=[
{
"Principal": "Arn",
"Actions": [
"quicksight:DescribeDataSource",
"quicksight:DescribeDataSourcePermissions",
"quicksight:PassDataSource",
],
}
],
)
]
assert mock_data_client.update_data_source.call_args_list == []
assert sorted(
mock_data_client.delete_data_source.call_args_list,
key=lambda x: x.kwargs["DataSourceId"],
) == [
mock.call(
AwsAccountId=mock.ANY,
DataSourceId="data-workspace-dev-my_database-88f3887d",
),
mock.call(
AwsAccountId=mock.ANY,
DataSourceId="data-workspace-dev-test_external_db2-88f3887d",
),
]
