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'