def test_sync_doesnt_create_role_if_user_already_has_role(
self, mock_hawk_request, create_tools_access_iam_role_task):
can_access_tools_permission = Permission.objects.get(
codename='start_all_applications',
content_type=ContentType.objects.get_for_model(
ApplicationInstance),
)
user = UserFactory.create(email='*****@*****.**')
user.user_permissions.add(can_access_tools_permission)
user.profile.sso_id = '00000000-0000-0000-0000-000000000000'
user.profile.tools_access_role_arn = 'some-arn'
user.save()
with open(
os.path.join(
os.path.dirname(__file__),
'test_fixture_activity_stream_sso_john_smith.json',
),
'r',
) as file:
user_john_smith = (200, file.read())
with open(
os.path.join(
os.path.dirname(__file__),
'test_fixture_activity_stream_sso_empty.json',
),
'r',
) as file:
empty_result = (200, file.read())
mock_hawk_request.side_effect = [user_john_smith, empty_result]
_do_sync_activity_stream_sso_users()
User = get_user_model()
all_users = User.objects.all()
assert len(all_users) == 1
assert not create_tools_access_iam_role_task.delay.called
def test_sync_updates_existing_users_sso_id_and_email(
self, mock_hawk_request):
# set the sso id to something different to what the activity stream
# will return and set the email to the third email in the list that
# the activity stream will return to test that it is able to look up
# the user and update both their email and sso id
user = UserFactory.create(email='*****@*****.**')
user.profile.sso_id = '00000000-0000-0000-0000-111111111111'
user.save()
with open(
os.path.join(
os.path.dirname(__file__),
'test_fixture_activity_stream_sso_john_smith_multiple_emails.json',
),
'r',
) as file:
user_john_smith = (200, file.read())
with open(
os.path.join(
os.path.dirname(__file__),
'test_fixture_activity_stream_sso_empty.json',
),
'r',
) as file:
empty_result = (200, file.read())
mock_hawk_request.side_effect = [user_john_smith, empty_result]
_do_sync_activity_stream_sso_users()
User = get_user_model()
all_users = User.objects.all()
assert len(all_users) == 1
assert (str(all_users[0].profile.sso_id) ==
'00000000-0000-0000-0000-000000000000')
assert str(all_users[0].email) == '*****@*****.**'
def test_metabase_link(self, mocker):
user = UserFactory.create()
vis = VisualisationCatalogueItemFactory.create(
user_access_type='REQUIRES_AUTHENTICATION')
link = VisualisationLinkFactory.create(
visualisation_type='METABASE',
identifier='123456789',
visualisation_catalogue_item=vis,
)
jwt_encode = mocker.patch(
'dataworkspace.apps.applications.views.jwt.encode')
jwt_encode.return_value = b'my-token'
client = Client(**get_http_sso_data(user))
response = client.get(link.get_absolute_url())
assert response.status_code == 200
assert (
'//metabase.dataworkspace.test:8000/embed/dashboard/my-token#bordered=false&titled=false'
in response.content.decode(response.charset))
assert ('frame-src metabase.dataworkspace.test'
in response['content-security-policy'])
def test_cleanup_temporary_query_tables(self, mock_connections,
mock_databases_data):
mock_cursor = Mock()
mock_connection = Mock()
mock_cursor_ctx_manager = MagicMock()
mock_cursor_ctx_manager.__enter__.return_value = mock_cursor
mock_connection.cursor.return_value = mock_cursor_ctx_manager
mock_connections.__getitem__.return_value = mock_connection
mock_databases_data.__getitem__.return_value = {
"USER": "******",
"NAME": "my_database"
"",
}
user = UserFactory()
user.profile.sso_id = '00000000-0000-0000-0000-000000000000' # yields a short hexdigest of 12b9377c
user.profile.save()
# last run 1 day and 1 hour ago so its materialized view should be deleted
with freeze_time(datetime.utcnow() - timedelta(days=1, hours=1)):
query_log_1 = QueryLogFactory.create(run_by_user=user)
# last run 2 hours ago so its materialized view should be kept
with freeze_time(datetime.utcnow() - timedelta(hours=2)):
QueryLogFactory.create(run_by_user=user)
cleanup_temporary_query_tables()
expected_calls = [
call('GRANT _user_12b9377c TO postgres'),
call(
f'DROP TABLE IF EXISTS _user_12b9377c._data_explorer_tmp_query_{query_log_1.id}'
),
call('REVOKE _user_12b9377c FROM postgres'),
]
mock_cursor.execute.assert_has_calls(expected_calls)
def test_cannot_open_playground_with_another_users_query(self, staff_client):
other_user = UserFactory(email="*****@*****.**")
query = SimpleQueryFactory(sql="select 1;", created_by_user=other_user)
resp = staff_client.get("%s?query_id=%s" % (reverse("explorer:index"), query.id))
assert resp.status_code == 404
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_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": "ADMIN",
"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
sync_quicksight_permissions(user_sso_ids_to_update=[
str(user.profile.sso_id),
str(user2.profile.sso_id)
])
# Assert
assert mock_user_client.update_user.call_args_list == [
mock.call(
AwsAccountId=mock.ANY,
Namespace='default',
Role='ADMIN',
UnapplyCustomPermissions=True,
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}',
),
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 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",
),
]
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=UserAccessType.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_custom_delimiter(self):
user = UserFactory()
q = SimpleQueryFactory(sql='select 1, 2')
exporter = CSVExporter(user=user, query=q)
res = exporter.get_output(delim='|')
assert res == '?column?|?column?\r\n1|2\r\n'
