milvus/internal/metastore/db/dao/user_role_test.go

package dao

import (
	"errors"
	"testing"

	"github.com/milvus-io/milvus/internal/metastore/db/dbmodel"

	"github.com/DATA-DOG/go-sqlmock"
	"github.com/stretchr/testify/assert"
)

func TestUserRole_GetUserRoles(t *testing.T) {
	var (
		userID1   = 1
		userID2   = 2
		roleID1   = 10
		roleID2   = 20
		userRoles []*dbmodel.UserRole
		err       error
	)

	// mock user and role
	mock.ExpectQuery("SELECT * FROM `user_role` WHERE `is_deleted` = ? AND `tenant_id` = ?").
		WithArgs(false, tenantID).
		WillReturnRows(
			sqlmock.NewRows([]string{"tenant_id", "user_id", "role_id"}).
				AddRow(tenantID, userID1, roleID1).
				AddRow(tenantID, userID2, roleID2))

	mock.ExpectQuery("SELECT * FROM `role` WHERE `role`.`id` IN (?,?)").
		WithArgs(roleID1, roleID2).
		WillReturnRows(
			sqlmock.NewRows([]string{"id", "tenant_id", "name"}).
				AddRow(roleID1, tenantID, "foo1").
				AddRow(roleID2, tenantID, "foo2"))

	mock.ExpectQuery("SELECT * FROM `credential_users` WHERE `credential_users`.`id` IN (?,?)").
		WithArgs(userID1, userID2).
		WillReturnRows(
			sqlmock.NewRows([]string{"id", "tenant_id", "username"}).
				AddRow(userID1, tenantID, "fo1").
				AddRow(userID2, tenantID, "fo2"))

	userRoles, err = userRoleTestDb.GetUserRoles(tenantID, 0, 0)
	mock.MatchExpectationsInOrder(false)
	assert.NoError(t, err)
	assert.Equal(t, 2, len(userRoles))
	assert.Equal(t, "foo1", userRoles[0].Role.Name)
	assert.Equal(t, "fo1", userRoles[0].User.Username)

	mock.ExpectQuery("SELECT * FROM `user_role` WHERE `is_deleted` = ? AND `tenant_id` = ?").
		WithArgs(false, tenantID).
		WillReturnError(errors.New("test error"))
	_, err = userRoleTestDb.GetUserRoles(tenantID, 0, 0)
	mock.MatchExpectationsInOrder(false)
	assert.Error(t, err)
}

func TestUserRole_GetUserRolesWithUserID(t *testing.T) {
	var (
		userID1   = 1
		roleID1   = 10
		roleID2   = 20
		userRoles []*dbmodel.UserRole
		err       error
	)

	mock.ExpectQuery("SELECT * FROM `user_role` WHERE `user_role`.`user_id` = ? AND `is_deleted` = ? AND `tenant_id` = ?").
		WithArgs(userID1, false, tenantID).
		WillReturnRows(
			sqlmock.NewRows([]string{"tenant_id", "user_id", "role_id"}).
				AddRow(tenantID, userID1, roleID1).
				AddRow(tenantID, userID1, roleID2))
	mock.ExpectQuery("SELECT * FROM `role` WHERE `role`.`id` IN (?,?)").
		WithArgs(roleID1, roleID2).
		WillReturnRows(
			sqlmock.NewRows([]string{"id", "tenant_id", "name"}).
				AddRow(roleID1, tenantID, "foo1").
				AddRow(roleID2, tenantID, "foo2"))
	mock.ExpectQuery("SELECT * FROM `credential_users` WHERE `credential_users`.`id` = ?").
		WithArgs(userID1).
		WillReturnRows(
			sqlmock.NewRows([]string{"id", "tenant_id", "username"}).
				AddRow(userID1, tenantID, "fo1"))

	userRoles, err = userRoleTestDb.GetUserRoles(tenantID, int64(userID1), 0)
	mock.MatchExpectationsInOrder(false)
	assert.NoError(t, err)
	assert.Equal(t, 2, len(userRoles))
	assert.Equal(t, "foo2", userRoles[1].Role.Name)
	assert.Equal(t, "fo1", userRoles[0].User.Username)
}

func TestUserRole_GetUserRolesWithRoleID(t *testing.T) {
	var (
		userID1   = 1
		userID2   = 2
		roleID1   = 10
		userRoles []*dbmodel.UserRole
		err       error
	)

	mock.ExpectQuery("SELECT * FROM `user_role` WHERE `user_role`.`role_id` = ? AND `is_deleted` = ? AND `tenant_id` = ?").
		WithArgs(roleID1, false, tenantID).
		WillReturnRows(
			sqlmock.NewRows([]string{"tenant_id", "user_id", "role_id"}).
				AddRow(tenantID, userID1, roleID1).
				AddRow(tenantID, userID2, roleID1))
	mock.ExpectQuery("SELECT * FROM `role` WHERE `role`.`id` = ?").
		WithArgs(roleID1).
		WillReturnRows(
			sqlmock.NewRows([]string{"id", "tenant_id", "name"}).
				AddRow(roleID1, tenantID, "foo1"))
	mock.ExpectQuery("SELECT * FROM `credential_users` WHERE `credential_users`.`id` IN (?,?)").
		WithArgs(userID1, userID2).
		WillReturnRows(
			sqlmock.NewRows([]string{"id", "tenant_id", "username"}).
				AddRow(userID1, tenantID, "fo1").
				AddRow(userID2, tenantID, "fo2"))

	userRoles, err = userRoleTestDb.GetUserRoles(tenantID, 0, int64(roleID1))
	mock.MatchExpectationsInOrder(false)
	assert.NoError(t, err)
	assert.Equal(t, 2, len(userRoles))
	assert.Equal(t, "foo1", userRoles[0].Role.Name)
	assert.Equal(t, "fo2", userRoles[1].User.Username)
}

func TestUserRole_Insert(t *testing.T) {
	var (
		userRole *dbmodel.UserRole
		err      error
	)
	userRole = &dbmodel.UserRole{
		Base:   GetBase(),
		UserID: 1,
		RoleID: 1,
	}

	mock.ExpectBegin()
	mock.ExpectExec("INSERT INTO `user_role` (`tenant_id`,`is_deleted`,`created_at`,`updated_at`,`user_id`,`role_id`) VALUES (?,?,?,?,?,?)").
		WithArgs(userRole.TenantID, userRole.IsDeleted, userRole.CreatedAt, userRole.UpdatedAt, userRole.UserID, userRole.RoleID).
		WillReturnResult(sqlmock.NewResult(1, 1))
	mock.ExpectCommit()
	err = userRoleTestDb.Insert(userRole)
	mock.MatchExpectationsInOrder(false)
	assert.NoError(t, err)
}

func TestUserRole_InsertError(t *testing.T) {
	var (
		userRole *dbmodel.UserRole
		err      error
	)
	userRole = &dbmodel.UserRole{
		Base:   GetBase(),
		UserID: 1,
		RoleID: 1,
	}

	mock.ExpectBegin()
	mock.ExpectExec("INSERT INTO `user_role` (`tenant_id`,`is_deleted`,`created_at`,`updated_at`,`user_id`,`role_id`) VALUES (?,?,?,?,?,?)").
		WithArgs(userRole.TenantID, userRole.IsDeleted, userRole.CreatedAt, userRole.UpdatedAt, userRole.UserID, userRole.RoleID).
		WillReturnError(errors.New("test error"))
	mock.ExpectRollback()
	err = userRoleTestDb.Insert(userRole)
	mock.MatchExpectationsInOrder(false)
	assert.Error(t, err)
}

func TestUserRole_Delete(t *testing.T) {
	var (
		userRole *dbmodel.UserRole
		getExec  func() *sqlmock.ExpectedExec
		err      error
	)
	userRole = &dbmodel.UserRole{
		Base:   GetBase(),
		UserID: 1,
		RoleID: 1,
	}
	getExec = func() *sqlmock.ExpectedExec {
		return mock.ExpectExec("UPDATE `user_role` SET `is_deleted`=?,`updated_at`=? WHERE `user_role`.`user_id` = ? AND `user_role`.`role_id` = ? AND `is_deleted` = ? AND `tenant_id` = ?").
			WithArgs(true, AnyTime{}, userRole.UserID, userRole.RoleID, userRole.IsDeleted, userRole.TenantID)
	}
	mock.ExpectBegin()
	getExec().WillReturnResult(sqlmock.NewResult(1, 1))
	mock.ExpectCommit()
	err = userRoleTestDb.Delete(userRole.TenantID, userRole.UserID, userRole.RoleID)
	mock.MatchExpectationsInOrder(false)
	assert.NoError(t, err)

	mock.ExpectBegin()
	getExec().WillReturnError(errors.New("test error"))
	mock.ExpectRollback()
	err = userRoleTestDb.Delete(userRole.TenantID, userRole.UserID, userRole.RoleID)
	mock.MatchExpectationsInOrder(false)
	assert.Error(t, err)
}