package role

import (
	"context"
	"errors"
	"testing"
	"time"

	"code.justin.tv/beefcake/server/internal/awsmocks"
	"code.justin.tv/beefcake/server/internal/config"
	"code.justin.tv/beefcake/server/internal/optional"
	"code.justin.tv/beefcake/server/internal/perm"
	"code.justin.tv/beefcake/server/internal/testconfig"
	"code.justin.tv/beefcake/server/rpc/beefcake"
	"github.com/aws/aws-sdk-go/aws"
	"github.com/aws/aws-sdk-go/aws/awserr"
	"github.com/aws/aws-sdk-go/service/dynamodb"
	"github.com/aws/aws-sdk-go/service/dynamodb/dynamodbattribute"
	"github.com/golang/protobuf/ptypes"
	"github.com/golang/protobuf/ptypes/timestamp"
	"github.com/stretchr/testify/assert"
	"github.com/stretchr/testify/mock"
	"github.com/stretchr/testify/require"
)

type rolesTest struct {
	Roles    *Roles
	Config   *config.Config
	DynamoDB *awsmocks.DynamoDBAPI
}

func newRolesTest(t *testing.T) *rolesTest {
	dynamoDB := new(awsmocks.DynamoDBAPI)
	config := testconfig.New(t)
	return &rolesTest{
		Roles: &Roles{
			Config:   config,
			DynamoDB: dynamoDB,
		},
		Config:   config,
		DynamoDB: dynamoDB,
	}
}

func (ct rolesTest) Teardown(t *testing.T) {
	ct.DynamoDB.AssertExpectations(t)
}

func TestRoles(t *testing.T) {
	const testRoleID = "test-role-id"
	const testRoleName = "test-role-name"

	testExpirationT := time.Now().Add(time.Hour).Truncate(time.Second)
	testExpiration := func(t *testing.T) *timestamp.Timestamp {
		out, err := ptypes.TimestampProto(testExpirationT)
		require.NoError(t, err)
		return out
	}

	testUserMembership := func(t *testing.T) *beefcake.Role_UserMembership {
		return &beefcake.Role_UserMembership{
			UserId:     "test-user-id",
			Expiration: testExpiration(t),
		}
	}

	testPermission := func() *beefcake.Permission {
		return &beefcake.Permission{Value: &beefcake.Permission_Legacy{
			Legacy: &beefcake.Permission_LegacyPermission{
				Id:   "p-id",
				Name: "p-name",
			},
		}}
	}

	marshalled := func(t *testing.T, in interface{}) *dynamodb.AttributeValue {
		out, err := dynamodbattribute.Marshal(in)
		require.NoError(t, err)
		return out
	}

	t.Run("All", func(t *testing.T) {
		testPermissions := func() perm.AttachedPermissions {
			return perm.AttachedPermissions([]*beefcake.AttachedPermission{
				{Id: "role-perm-id", Permission: testPermission()},
			})
		}

		marshalledRole := func(t *testing.T, r *Role) map[string]*dynamodb.AttributeValue {
			m, err := dynamodbattribute.MarshalMap(r)
			require.NoError(t, err)
			return m
		}

		tcs := []struct {
			Name         string
			Last         bool
			ExpectedCont bool
		}{
			{Name: "has more", Last: false, ExpectedCont: true},
			{Name: "last page", Last: true, ExpectedCont: false},
		}

		for _, tc := range tcs {
			t.Run("Success "+tc.Name, func(t *testing.T) {
				ct := newRolesTest(t)
				defer ct.Teardown(t)

				ct.DynamoDB.
					On("ScanPagesWithContext", mock.Anything, &dynamodb.ScanInput{
						Limit:     aws.Int64(100),
						TableName: aws.String(ct.Config.RolesTableName.Get()),
					}, mock.Anything).
					Run(func(args mock.Arguments) {
						fn := args.Get(2).(func(*dynamodb.ScanOutput, bool) bool)
						assert.Equal(t, tc.ExpectedCont, fn(&dynamodb.ScanOutput{
							Items: []map[string]*dynamodb.AttributeValue{
								marshalledRole(t, &Role{ID: "role-id-2", Permissions: testPermissions(), UserMemberships: UserMemberships{}}),
								marshalledRole(t, &Role{ID: "role-id-3", Permissions: testPermissions(), UserMemberships: UserMemberships{}}),
								marshalledRole(t, &Role{ID: "role-id-1", Permissions: testPermissions(), UserMemberships: UserMemberships{}}),
							},
						}, tc.Last))
					}).
					Return(nil)

				res, err := ct.Roles.All(context.Background())
				require.NoError(t, err)
				assert.Equal(t, []*Role{
					{ID: "role-id-1", Permissions: testPermissions(), UserMemberships: UserMemberships{}},
					{ID: "role-id-2", Permissions: testPermissions(), UserMemberships: UserMemberships{}},
					{ID: "role-id-3", Permissions: testPermissions(), UserMemberships: UserMemberships{}},
				}, res)
			})
		}
	})

	t.Run("Create", func(t *testing.T) {
		const testRoleName = "test-role-name"

		t.Run("Success", func(t *testing.T) {
			ct := newRolesTest(t)
			defer ct.Teardown(t)

			var roleID string
			ct.DynamoDB.
				On("PutItemWithContext", mock.Anything, mock.Anything).
				Run(func(args mock.Arguments) {
					input := args.Get(1).(*dynamodb.PutItemInput)
					roleID = aws.StringValue(input.Item["ID"].S)

					assert.Equal(t, &dynamodb.PutItemInput{
						ConditionExpression: aws.String("attribute_not_exists (#0)"),
						ExpressionAttributeNames: map[string]*string{
							"#0": aws.String(ct.Config.RolesHashKey.Get()),
						},
						Item: map[string]*dynamodb.AttributeValue{
							"ID":              {S: aws.String(roleID)},
							"Name":            {S: aws.String(testRoleName)},
							"Permissions":     {M: map[string]*dynamodb.AttributeValue{}},
							"UserMemberships": {M: map[string]*dynamodb.AttributeValue{}},
						},
						TableName: aws.String(ct.Config.RolesTableName.Get()),
					}, input)
				}).
				Return(&dynamodb.PutItemOutput{}, nil)

			res, err := ct.Roles.Create(context.Background(), testRoleName)
			require.NoError(t, err)
			assert.Equal(t, &Role{ID: roleID, Name: testRoleName}, res)
		})

		t.Run("Empty role name should fail", func(t *testing.T) {
			ct := newRolesTest(t)
			defer ct.Teardown(t)

			_, err := ct.Roles.Create(context.Background(), "")
			assert.Equal(t, ErrMissingRoleParameters{Parameters: []string{"Name"}}, err)
		})
	})

	t.Run("Get", func(t *testing.T) {
		testRole := func(t *testing.T) *Role {
			return &Role{
				ID:   testRoleID,
				Name: "test-name",
				Permissions: perm.AttachedPermissions([]*beefcake.AttachedPermission{
					{Id: "role-perm-id", Permission: testPermission()},
				}),
				UserMemberships: UserMemberships([]*beefcake.Role_UserMembership{
					testUserMembership(t),
				}),
			}
		}

		t.Run("Success", func(t *testing.T) {
			ct := newRolesTest(t)
			defer ct.Teardown(t)

			item, err := dynamodbattribute.MarshalMap(testRole(t))
			require.NoError(t, err)

			ct.DynamoDB.
				On("GetItemWithContext", mock.Anything, &dynamodb.GetItemInput{
					Key: map[string]*dynamodb.AttributeValue{
						ct.Config.RolesHashKey.Get(): {S: aws.String(testRoleID)},
					},
					TableName: aws.String(ct.Config.RolesTableName.Get()),
				}).
				Return(&dynamodb.GetItemOutput{Item: item}, nil)

			res, err := ct.Roles.Get(context.Background(), testRoleID)
			require.NoError(t, err)

			assert.Equal(t, testExpiration(t), res.UserMemberships[0].Expiration)
			assert.Equal(t, testRole(t), res)
		})
	})

	t.Run("AddPermission", func(t *testing.T) {
		t.Run("Success", func(t *testing.T) {
			ct := newRolesTest(t)
			defer ct.Teardown(t)

			testPermission := func() *beefcake.Permission {
				return &beefcake.Permission{Value: &beefcake.Permission_TwitchUser{
					TwitchUser: &beefcake.Permission_TwitchUserPermission{},
				}}
			}

			ct.DynamoDB.
				On("UpdateItemWithContext", mock.Anything, mock.Anything).
				Run(func(arg mock.Arguments) {
					input := arg.Get(1).(*dynamodb.UpdateItemInput)

					// rolePermID is randomly generated
					rolePermID := aws.StringValue(input.ExpressionAttributeNames["#2"])

					expectedRolePerms, err := dynamodbattribute.Marshal(perm.Permission(*testPermission()))
					require.NoError(t, err)

					assert.Equal(t, &dynamodb.UpdateItemInput{
						ConditionExpression: aws.String("attribute_exists (#0)"),
						ExpressionAttributeNames: map[string]*string{
							"#0": aws.String(ct.Config.RolesHashKey.Get()),
							"#1": aws.String(PermissionsAttribute),
							"#2": aws.String(rolePermID),
						},
						ExpressionAttributeValues: map[string]*dynamodb.AttributeValue{
							":0": expectedRolePerms,
						},
						Key: map[string]*dynamodb.AttributeValue{
							ct.Config.RolesHashKey.Get(): {S: aws.String(testRoleID)},
						},
						TableName:        aws.String(ct.Config.RolesTableName.Get()),
						UpdateExpression: aws.String("SET #1.#2 = :0\n"),
					}, input)
				}).
				Return(nil, nil)

			require.NoError(t, ct.Roles.AddPermission(context.Background(), testRoleID, testPermission()))
		})

		t.Run("Legacy permission should fail", func(t *testing.T) {
			ct := newRolesTest(t)
			defer ct.Teardown(t)

			testPermission := func() *beefcake.Permission {
				return &beefcake.Permission{Value: &beefcake.Permission_Legacy{
					Legacy: &beefcake.Permission_LegacyPermission{
						Id:   "p-id",
						Name: "p-name",
					},
				}}
			}

			assert.Equal(t, ErrCannotAddLegacyPermission, ct.Roles.AddPermission(context.Background(), testRoleID, testPermission()))
		})
	})

	t.Run("Update", func(t *testing.T) {
		const testLegacyPermissionID = "test-legacy-permission-id"

		mockUpdate := func(ct *rolesTest) *mock.Call {
			return ct.DynamoDB.
				On("UpdateItemWithContext", mock.Anything, &dynamodb.UpdateItemInput{
					ConditionExpression: aws.String("attribute_exists (#0)"),
					ExpressionAttributeNames: map[string]*string{
						"#0": aws.String(ct.Config.RolesHashKey.Get()),
						"#1": aws.String(NameAttribute),
						"#2": aws.String(LegacyPermissionsAttribute),
					},
					ExpressionAttributeValues: map[string]*dynamodb.AttributeValue{
						":0": {S: aws.String(testRoleName)},
						":1": {SS: []*string{aws.String(testLegacyPermissionID)}},
					},
					Key: map[string]*dynamodb.AttributeValue{
						ct.Config.RolesHashKey.Get(): {S: aws.String(testRoleID)},
					},
					TableName:        aws.String(ct.Config.RolesTableName.Get()),
					UpdateExpression: aws.String("SET #1 = :0, #2 = :1\n"),
				})
		}

		t.Run("success", func(t *testing.T) {
			ct := newRolesTest(t)
			defer ct.Teardown(t)

			mockUpdate(ct).
				Return(nil, nil)

			require.NoError(t, ct.Roles.Update(
				context.Background(),
				testRoleID,
				UpdateOptions{
					Name:                optional.String(testRoleName),
					LegacyPermissionIDs: optional.StringSlice([]string{testLegacyPermissionID}),
				}))
		})

		t.Run("missing should return ErrRoleDoesNotExist", func(t *testing.T) {
			ct := newRolesTest(t)
			defer ct.Teardown(t)

			mockUpdate(ct).
				Return(nil, awserr.New(dynamodb.ErrCodeConditionalCheckFailedException, "", errors.New("")))

			assert.Equal(t, ErrRoleDoesNotExist{RoleID: testRoleID}, ct.Roles.Update(
				context.Background(),
				testRoleID,
				UpdateOptions{
					Name:                optional.String(testRoleName),
					LegacyPermissionIDs: optional.StringSlice([]string{testLegacyPermissionID}),
				}))
		})

		t.Run("success - no name", func(t *testing.T) {
			ct := newRolesTest(t)
			defer ct.Teardown(t)

			ct.DynamoDB.
				On("UpdateItemWithContext", mock.Anything, &dynamodb.UpdateItemInput{
					ConditionExpression: aws.String("attribute_exists (#0)"),
					ExpressionAttributeNames: map[string]*string{
						"#0": aws.String(ct.Config.RolesHashKey.Get()),
						"#1": aws.String(LegacyPermissionsAttribute),
					},
					ExpressionAttributeValues: map[string]*dynamodb.AttributeValue{
						":0": {SS: []*string{aws.String(testLegacyPermissionID)}},
					},
					Key: map[string]*dynamodb.AttributeValue{
						ct.Config.RolesHashKey.Get(): {S: aws.String(testRoleID)},
					},
					TableName:        aws.String(ct.Config.RolesTableName.Get()),
					UpdateExpression: aws.String("SET #1 = :0\n"),
				}).
				Return(nil, nil)

			require.NoError(t, ct.Roles.Update(
				context.Background(),
				testRoleID,
				UpdateOptions{
					LegacyPermissionIDs: optional.StringSlice([]string{testLegacyPermissionID}),
				}))
		})

		t.Run("success - no legacy", func(t *testing.T) {
			ct := newRolesTest(t)
			defer ct.Teardown(t)

			ct.DynamoDB.
				On("UpdateItemWithContext", mock.Anything, &dynamodb.UpdateItemInput{
					ConditionExpression: aws.String("attribute_exists (#0)"),
					ExpressionAttributeNames: map[string]*string{
						"#0": aws.String(ct.Config.RolesHashKey.Get()),
						"#1": aws.String(NameAttribute),
					},
					ExpressionAttributeValues: map[string]*dynamodb.AttributeValue{
						":0": {S: aws.String(testRoleName)},
					},
					Key: map[string]*dynamodb.AttributeValue{
						ct.Config.RolesHashKey.Get(): {S: aws.String(testRoleID)},
					},
					TableName:        aws.String(ct.Config.RolesTableName.Get()),
					UpdateExpression: aws.String("SET #1 = :0\n"),
				}).
				Return(nil, nil)

			require.NoError(t, ct.Roles.Update(
				context.Background(),
				testRoleID,
				UpdateOptions{
					Name: optional.String(testRoleName),
				}))
		})

		t.Run("success - no updates", func(t *testing.T) {
			ct := newRolesTest(t)
			defer ct.Teardown(t)

			require.NoError(t, ct.Roles.Update(
				context.Background(),
				testRoleID,
				UpdateOptions{}))
		})

		t.Run("permissions of len 0 should delete the attribute", func(t *testing.T) {
			ct := newRolesTest(t)
			defer ct.Teardown(t)

			ct.DynamoDB.
				On("UpdateItemWithContext", mock.Anything, &dynamodb.UpdateItemInput{
					ConditionExpression: aws.String("attribute_exists (#0)"),
					ExpressionAttributeNames: map[string]*string{
						"#0": aws.String(ct.Config.RolesHashKey.Get()),
						"#1": aws.String(LegacyPermissionsAttribute),
						"#2": aws.String(NameAttribute),
					},
					ExpressionAttributeValues: map[string]*dynamodb.AttributeValue{
						":0": {S: aws.String(testRoleName)},
					},
					Key: map[string]*dynamodb.AttributeValue{
						ct.Config.RolesHashKey.Get(): {S: aws.String(testRoleID)},
					},
					TableName:        aws.String(ct.Config.RolesTableName.Get()),
					UpdateExpression: aws.String("REMOVE #1\nSET #2 = :0\n"),
				}).
				Return(nil, nil)

			require.NoError(t, ct.Roles.Update(
				context.Background(),
				testRoleID,
				UpdateOptions{
					Name:                optional.String(testRoleName),
					LegacyPermissionIDs: optional.StringSlice([]string{}),
				}))
		})
	})

	t.Run("AddLegacyPermission", func(t *testing.T) {
		const testPermissionID = "test-permission-id"

		legacyPermissionCheck := func(ct *rolesTest) *dynamodb.TransactWriteItem {
			return &dynamodb.TransactWriteItem{
				ConditionCheck: &dynamodb.ConditionCheck{
					ConditionExpression: aws.String("attribute_exists (#0)"),
					ExpressionAttributeNames: map[string]*string{
						"#0": aws.String(ct.Config.LegacyPermissionsHashKey.Get()),
					},
					Key: map[string]*dynamodb.AttributeValue{
						ct.Config.LegacyPermissionsHashKey.Get(): {S: aws.String(testPermissionID)},
					},
					TableName: aws.String(ct.Config.LegacyPermissionsTableName.Get()),
				},
			}
		}

		setCall := func(t *testing.T, ct *rolesTest) *mock.Call {
			return ct.DynamoDB.
				On("TransactWriteItemsWithContext", mock.Anything, &dynamodb.TransactWriteItemsInput{
					TransactItems: []*dynamodb.TransactWriteItem{
						legacyPermissionCheck(ct),
						{
							Update: &dynamodb.Update{
								ConditionExpression: aws.String("attribute_not_exists (#0)"),
								ExpressionAttributeNames: map[string]*string{
									"#0": aws.String(LegacyPermissionsAttribute),
								},
								ExpressionAttributeValues: map[string]*dynamodb.AttributeValue{
									":0": {SS: []*string{aws.String(testPermissionID)}},
								},
								Key: map[string]*dynamodb.AttributeValue{
									ct.Config.RolesHashKey.Get(): {S: aws.String(testRoleID)},
								},
								TableName:        aws.String(ct.Config.RolesTableName.Get()),
								UpdateExpression: aws.String("SET #0 = :0\n"),
							},
						},
					},
				})
		}

		updateCall := func(t *testing.T, ct *rolesTest) *mock.Call {
			return ct.DynamoDB.
				On("TransactWriteItemsWithContext", mock.Anything, &dynamodb.TransactWriteItemsInput{
					TransactItems: []*dynamodb.TransactWriteItem{
						legacyPermissionCheck(ct),
						{
							Update: &dynamodb.Update{
								ConditionExpression: aws.String("attribute_exists (#0)"),
								ExpressionAttributeNames: map[string]*string{
									"#0": aws.String(LegacyPermissionsAttribute),
								},
								ExpressionAttributeValues: map[string]*dynamodb.AttributeValue{
									":0": {SS: []*string{aws.String(testPermissionID)}},
								},
								Key: map[string]*dynamodb.AttributeValue{
									ct.Config.RolesHashKey.Get(): {S: aws.String(testRoleID)},
								},
								TableName:        aws.String(ct.Config.RolesTableName.Get()),
								UpdateExpression: aws.String("ADD #0 :0\n"),
							},
						},
					},
				})
		}

		t.Run("Success - set", func(t *testing.T) {
			ct := newRolesTest(t)
			defer ct.Teardown(t)

			setCall(t, ct).
				Return(nil, nil)

			require.NoError(t, ct.Roles.AddLegacyPermission(context.Background(), testRoleID, testPermissionID))
		})

		t.Run("Success - update", func(t *testing.T) {
			ct := newRolesTest(t)
			defer ct.Teardown(t)

			setCall(t, ct).
				Return(nil, awserr.New(dynamodb.ErrCodeTransactionCanceledException, "", errors.New("")))
			updateCall(t, ct).
				Return(nil, nil)

			require.NoError(t, ct.Roles.AddLegacyPermission(context.Background(), testRoleID, testPermissionID))
		})
	})

	t.Run("RemoveLegacyPermission", func(t *testing.T) {
		const testLegacyPermissionID = "test-legacy-permission-id"

		t.Run("Success", func(t *testing.T) {
			ct := newRolesTest(t)
			defer ct.Teardown(t)

			ct.DynamoDB.
				On("UpdateItemWithContext", mock.Anything, &dynamodb.UpdateItemInput{
					ConditionExpression: aws.String("attribute_exists (#0)"),
					ExpressionAttributeNames: map[string]*string{
						"#0": aws.String(LegacyPermissionsAttribute),
					},
					ExpressionAttributeValues: map[string]*dynamodb.AttributeValue{
						":0": {
							SS: []*string{aws.String(testLegacyPermissionID)},
						},
					},
					Key: map[string]*dynamodb.AttributeValue{
						ct.Config.RolesHashKey.Get(): {S: aws.String(testRoleID)},
					},
					TableName:        aws.String(ct.Config.RolesTableName.Get()),
					UpdateExpression: aws.String("DELETE #0 :0\n"),
				}).
				Return(nil, nil)

			require.NoError(t, ct.Roles.RemoveLegacyPermission(
				context.Background(),
				testRoleID,
				testLegacyPermissionID))
		})
	})

	t.Run("AddLegacyPermissionDetails", func(t *testing.T) {
		const testLegacyPermissionID = "test-legacy-permission-id"
		const testLegacyPermissionName = "test-legacy-permission-name"

		t.Run("Success", func(t *testing.T) {
			ct := newRolesTest(t)
			defer ct.Teardown(t)

			ct.DynamoDB.
				On("UpdateItemWithContext", mock.Anything, &dynamodb.UpdateItemInput{
					ConditionExpression: aws.String("attribute_exists (#0)"),
					ExpressionAttributeNames: map[string]*string{
						"#0": aws.String(ct.Config.RolesHashKey.Get()),
						"#1": aws.String(PermissionsAttribute),
						"#2": aws.String("legacy:" + testLegacyPermissionID),
					},
					ExpressionAttributeValues: map[string]*dynamodb.AttributeValue{
						":0": marshalled(t, perm.Permission(beefcake.Permission{
							Value: &beefcake.Permission_Legacy{
								Legacy: &beefcake.Permission_LegacyPermission{
									Id:   testLegacyPermissionID,
									Name: testLegacyPermissionName,
								},
							},
						})),
					},
					Key: map[string]*dynamodb.AttributeValue{
						ct.Config.RolesHashKey.Get(): {S: aws.String(testRoleID)},
					},
					TableName:        aws.String(ct.Config.RolesTableName.Get()),
					UpdateExpression: aws.String("SET #1.#2 = :0\n"),
				}).
				Return(nil, nil)

			require.NoError(t, ct.Roles.AddLegacyPermissionDetails(
				context.Background(),
				testRoleID,
				&beefcake.Permission_LegacyPermission{
					Id:   testLegacyPermissionID,
					Name: testLegacyPermissionName,
				}))
		})
	})

	t.Run("RemoveLegacyPermissionDetails", func(t *testing.T) {
		const testLegacyPermissionID = "test-legacy-permission-id"

		t.Run("Success", func(t *testing.T) {
			ct := newRolesTest(t)
			defer ct.Teardown(t)

			ct.DynamoDB.
				On("UpdateItemWithContext", mock.Anything, &dynamodb.UpdateItemInput{
					ConditionExpression: aws.String("attribute_exists (#0)"),
					ExpressionAttributeNames: map[string]*string{
						"#0": aws.String(ct.Config.RolesHashKey.Get()),
						"#1": aws.String(PermissionsAttribute),
						"#2": aws.String("legacy:" + testLegacyPermissionID),
					},
					Key: map[string]*dynamodb.AttributeValue{
						ct.Config.RolesHashKey.Get(): {S: aws.String(testRoleID)},
					},
					TableName:        aws.String(ct.Config.RolesTableName.Get()),
					UpdateExpression: aws.String("REMOVE #1.#2\n"),
				}).
				Return(nil, nil)

			require.NoError(t, ct.Roles.RemoveLegacyPermissionDetails(
				context.Background(),
				testRoleID,
				testLegacyPermissionID,
			))
		})
	})

	t.Run("Override", func(t *testing.T) {
		const legacyPermID = "p-id"

		legacyPerm := func() *beefcake.Permission {
			return &beefcake.Permission{Value: &beefcake.Permission_Legacy{
				Legacy: &beefcake.Permission_LegacyPermission{
					Id:   legacyPermID,
					Name: "p-name",
				},
			}}
		}

		otherPerm := func() *beefcake.Permission {
			return &beefcake.Permission{Value: &beefcake.Permission_TwitchUser{
				TwitchUser: &beefcake.Permission_TwitchUserPermission{
					Action: beefcake.Permission_TwitchUserPermission_ANY,
				},
			}}
		}

		t.Run("Success", func(t *testing.T) {
			ct := newRolesTest(t)
			defer ct.Teardown(t)

			ct.DynamoDB.
				On("UpdateItemWithContext", mock.Anything, mock.Anything).
				Run(func(arg mock.Arguments) {
					input := arg.Get(1).(*dynamodb.UpdateItemInput)

					// rolePermID is randomly generated
					rolePermissions := input.ExpressionAttributeValues[":1"].M
					require.Len(t, rolePermissions, 1)
					var rolePermID string
					for key := range rolePermissions {
						rolePermID = key
					}

					require.Equal(t, map[string]*dynamodb.AttributeValue{
						":0": marshalled(t, testRoleName),
						":1": marshalled(t, perm.AttachedPermissions([]*beefcake.AttachedPermission{
							{Id: rolePermID, Permission: otherPerm()},
						})),
						":2": marshalled(t, UserMemberships([]*beefcake.Role_UserMembership{
							testUserMembership(t),
						})),
						":3": {SS: []*string{
							aws.String(legacyPermID),
						}},
					}, input.ExpressionAttributeValues)

					assert.Equal(t, &dynamodb.UpdateItemInput{
						ExpressionAttributeNames: map[string]*string{
							"#0": aws.String(NameAttribute),
							"#1": aws.String(PermissionsAttribute),
							"#2": aws.String(UserMembershipsAttribute),
							"#3": aws.String(LegacyPermissionsAttribute),
						},
						ExpressionAttributeValues: map[string]*dynamodb.AttributeValue{
							":0": marshalled(t, testRoleName),
							":1": marshalled(t, perm.AttachedPermissions([]*beefcake.AttachedPermission{
								{Id: rolePermID, Permission: otherPerm()},
							})),
							":2": marshalled(t, UserMemberships([]*beefcake.Role_UserMembership{
								testUserMembership(t),
							})),
							":3": {SS: []*string{
								aws.String(legacyPermID),
							}},
						},
						Key: map[string]*dynamodb.AttributeValue{
							ct.Config.RolesHashKey.Get(): {S: aws.String(testRoleID)},
						},
						TableName:        aws.String(ct.Config.RolesTableName.Get()),
						UpdateExpression: aws.String("SET #0 = :0, #1 = :1, #2 = :2, #3 = :3\n"),
					}, input)
				}).
				Return(nil, nil)

			require.NoError(t, ct.Roles.Override(
				context.Background(),
				testRoleID,
				testRoleName,
				[]*beefcake.Permission{
					legacyPerm(),
					otherPerm(),
				},
				[]*beefcake.Role_UserMembership{
					testUserMembership(t),
				},
			))
		})

		t.Run("No legacy permissions", func(t *testing.T) {
			ct := newRolesTest(t)
			defer ct.Teardown(t)

			ct.DynamoDB.
				On("UpdateItemWithContext", mock.Anything, mock.Anything).
				Run(func(arg mock.Arguments) {
					input := arg.Get(1).(*dynamodb.UpdateItemInput)

					// rolePermID is randomly generated
					rolePermissions := input.ExpressionAttributeValues[":1"].M
					require.Len(t, rolePermissions, 1)
					var rolePermID string
					for key := range rolePermissions {
						rolePermID = key
					}

					assert.Equal(t, &dynamodb.UpdateItemInput{
						ExpressionAttributeNames: map[string]*string{
							"#0": aws.String(NameAttribute),
							"#1": aws.String(PermissionsAttribute),
							"#2": aws.String(UserMembershipsAttribute),
						},
						ExpressionAttributeValues: map[string]*dynamodb.AttributeValue{
							":0": marshalled(t, testRoleName),
							":1": marshalled(t, perm.AttachedPermissions([]*beefcake.AttachedPermission{
								{Id: rolePermID, Permission: otherPerm()},
							})),
							":2": marshalled(t, UserMemberships([]*beefcake.Role_UserMembership{
								testUserMembership(t),
							})),
						},
						Key: map[string]*dynamodb.AttributeValue{
							ct.Config.RolesHashKey.Get(): {S: aws.String(testRoleID)},
						},
						TableName:        aws.String(ct.Config.RolesTableName.Get()),
						UpdateExpression: aws.String("SET #0 = :0, #1 = :1, #2 = :2\n"),
					}, input)
				}).
				Return(nil, nil)

			require.NoError(t, ct.Roles.Override(
				context.Background(),
				testRoleID,
				testRoleName,
				[]*beefcake.Permission{
					otherPerm(),
				},
				[]*beefcake.Role_UserMembership{
					testUserMembership(t),
				},
			))
		})
	})

	t.Run("RemovePermission", func(t *testing.T) {
		const testRolePermissionID = "test-role-permission-id"

		mockRemoveBothCall := func(ct *rolesTest) *mock.Call {
			return ct.DynamoDB.
				On("UpdateItemWithContext", mock.Anything, &dynamodb.UpdateItemInput{
					ConditionExpression: aws.String("attribute_exists (#0)"),
					ExpressionAttributeNames: map[string]*string{
						"#0": aws.String(LegacyPermissionsAttribute),
						"#1": aws.String(PermissionsAttribute),
						"#2": aws.String(testRolePermissionID),
					},
					ExpressionAttributeValues: map[string]*dynamodb.AttributeValue{
						":0": {
							SS: []*string{aws.String(testRolePermissionID)},
						},
					},
					Key: map[string]*dynamodb.AttributeValue{
						ct.Config.RolesHashKey.Get(): {S: aws.String(testRoleID)},
					},
					TableName:        aws.String(ct.Config.RolesTableName.Get()),
					UpdateExpression: aws.String("DELETE #0 :0\nREMOVE #1.#2\n"),
				})
		}

		mockRemoveOneCall := func(ct *rolesTest) *mock.Call {
			return ct.DynamoDB.
				On("UpdateItemWithContext", mock.Anything, &dynamodb.UpdateItemInput{
					ConditionExpression: aws.String("attribute_not_exists (#0)"),
					ExpressionAttributeNames: map[string]*string{
						"#0": aws.String(LegacyPermissionsAttribute),
						"#1": aws.String(PermissionsAttribute),
						"#2": aws.String(testRolePermissionID),
					},
					Key: map[string]*dynamodb.AttributeValue{
						ct.Config.RolesHashKey.Get(): {S: aws.String(testRoleID)},
					},
					TableName:        aws.String(ct.Config.RolesTableName.Get()),
					UpdateExpression: aws.String("REMOVE #1.#2\n"),
				})
		}

		t.Run("Success with legacy entry", func(t *testing.T) {
			ct := newRolesTest(t)
			defer ct.Teardown(t)

			mockRemoveBothCall(ct).
				Return(nil, nil)

			require.NoError(t, ct.Roles.RemovePermission(context.Background(), testRoleID, testRolePermissionID))
		})

		t.Run("Success with one entry", func(t *testing.T) {
			ct := newRolesTest(t)
			defer ct.Teardown(t)

			mockRemoveBothCall(ct).
				Return(nil, awserr.New(dynamodb.ErrCodeConditionalCheckFailedException, "", errors.New("")))

			mockRemoveOneCall(ct).
				Return(nil, nil)

			require.NoError(t, ct.Roles.RemovePermission(context.Background(), testRoleID, testRolePermissionID))
		})
	})

	t.Run("AddUserMembership", func(t *testing.T) {
		t.Run("Success", func(t *testing.T) {
			ct := newRolesTest(t)
			defer ct.Teardown(t)

			ct.DynamoDB.
				On("UpdateItemWithContext", mock.Anything, &dynamodb.UpdateItemInput{
					ExpressionAttributeNames: map[string]*string{
						"#0": aws.String(UserMembershipsAttribute),
						"#1": aws.String(testUserMembership(t).GetUserId()),
					},
					ExpressionAttributeValues: map[string]*dynamodb.AttributeValue{
						":0": marshalled(t, UserMembership(*testUserMembership(t))),
					},
					Key: map[string]*dynamodb.AttributeValue{
						ct.Config.RolesHashKey.Get(): {S: aws.String(testRoleID)},
					},
					TableName:        aws.String(ct.Config.RolesTableName.Get()),
					UpdateExpression: aws.String("SET #0.#1 = :0\n"),
				}).
				Return(nil, nil)

			require.NoError(t, ct.Roles.AddUserMembership(context.Background(), testRoleID, testUserMembership(t)))
		})
	})

	t.Run("RemoveUserMembership", func(t *testing.T) {
		const testUserID = "test-user-id"

		t.Run("Success", func(t *testing.T) {
			ct := newRolesTest(t)
			defer ct.Teardown(t)

			ct.DynamoDB.
				On("UpdateItemWithContext", mock.Anything, &dynamodb.UpdateItemInput{
					ExpressionAttributeNames: map[string]*string{
						"#0": aws.String(UserMembershipsAttribute),
						"#1": aws.String(testUserID),
					},
					Key: map[string]*dynamodb.AttributeValue{
						ct.Config.RolesHashKey.Get(): {S: aws.String(testRoleID)},
					},
					TableName:        aws.String(ct.Config.RolesTableName.Get()),
					UpdateExpression: aws.String("REMOVE #0.#1\n"),
				}).
				Return(nil, nil)

			require.NoError(t, ct.Roles.RemoveUserMembership(context.Background(), testRoleID, testUserID))
		})
	})

	t.Run("Delete", func(t *testing.T) {
		t.Run("Success", func(t *testing.T) {
			ct := newRolesTest(t)
			defer ct.Teardown(t)

			ct.DynamoDB.
				On("DeleteItemWithContext", mock.Anything, &dynamodb.DeleteItemInput{
					Key: map[string]*dynamodb.AttributeValue{
						ct.Config.RolesHashKey.Get(): {S: aws.String(testRoleID)},
					},
					TableName: aws.String(ct.Config.RolesTableName.Get()),
				}).
				Return(nil, nil)

			require.NoError(t, ct.Roles.Delete(context.Background(), testRoleID))
		})
	})
}
