add ut

Signed-off-by: shaoting-huang <[email protected]>
shaoting-huang · Nov 15, 2024 · 04ee212 · 04ee212
1 parent 10c60c2
commit 04ee212
Show file tree

Hide file tree

Showing 2 changed files with 25 additions and 5 deletions.
diff --git a/internal/rootcoord/root_coord.go b/internal/rootcoord/root_coord.go
@@ -630,7 +630,7 @@ func (c *Core) initPublicRolePrivilege() error {
 }
 
 func (c *Core) initBuiltinPrivilegeGroups() error {
-	// init built in privilege groups, if config enabled, override by config
+	// init built in privilege groups, override by config if rbac config enabled
 	for groupName, privileges := range util.BuiltinPrivilegeGroups {
 		if err := c.meta.CreatePrivilegeGroup(groupName); err != nil {
 			return err
@@ -639,10 +639,7 @@ func (c *Core) initBuiltinPrivilegeGroups() error {
 			var confPrivs []string
 			switch groupName {
 			case "ClusterReadOnly":
-				confPrivs := Params.RbacConfig.ClusterReadOnlyPrivileges.GetAsStrings()
-				if len(confPrivs) > 0 {
-					privileges = confPrivs
-				}
+				confPrivs = Params.RbacConfig.ClusterReadOnlyPrivileges.GetAsStrings()
 			case "ClusterReadWrite":
 				confPrivs = Params.RbacConfig.ClusterReadWritePrivileges.GetAsStrings()
 			case "ClusterAdmin":

diff --git a/internal/rootcoord/root_coord_test.go b/internal/rootcoord/root_coord_test.go
@@ -1976,6 +1976,8 @@ func TestCore_InitRBAC(t *testing.T) {
 		c := newTestCore(withHealthyCode(), withMeta(meta))
 		meta.EXPECT().CreateRole(mock.Anything, mock.Anything).Return(nil).Twice()
 		meta.EXPECT().OperatePrivilege(mock.Anything, mock.Anything, mock.Anything).Return(nil).Twice()
+		meta.EXPECT().CreatePrivilegeGroup(mock.Anything).Return(nil).Times(len(util.BuiltinPrivilegeGroups))
+		meta.EXPECT().OperatePrivilegeGroup(mock.Anything, mock.Anything, mock.Anything).Return(nil).Times(len(util.BuiltinPrivilegeGroups))
 
 		Params.Save(Params.RoleCfg.Enabled.Key, "false")
 		Params.Save(Params.ProxyCfg.EnablePublicPrivilege.Key, "true")
@@ -1995,6 +1997,8 @@ func TestCore_InitRBAC(t *testing.T) {
 		c := newTestCore(withHealthyCode(), withMeta(meta))
 		meta.EXPECT().CreateRole(mock.Anything, mock.Anything).Return(nil).Times(3)
 		meta.EXPECT().OperatePrivilege(mock.Anything, mock.Anything, mock.Anything).Return(nil).Once()
+		meta.EXPECT().CreatePrivilegeGroup(mock.Anything).Return(nil).Times(len(util.BuiltinPrivilegeGroups))
+		meta.EXPECT().OperatePrivilegeGroup(mock.Anything, mock.Anything, mock.Anything).Return(nil).Times(len(util.BuiltinPrivilegeGroups))
 
 		Params.Save(Params.RoleCfg.Enabled.Key, "true")
 		Params.Save(Params.RoleCfg.Roles.Key, builtinRoles)
@@ -2009,6 +2013,25 @@ func TestCore_InitRBAC(t *testing.T) {
 		err := c.initRbac()
 		assert.NoError(t, err)
 	})
+
+	t.Run("init default privilege groups", func(t *testing.T) {
+		clusterReadWrite := `SelectOwnership,SelectUser,DescribeResourceGroup`
+		meta := mockrootcoord.NewIMetaTable(t)
+		c := newTestCore(withHealthyCode(), withMeta(meta))
+		meta.EXPECT().CreatePrivilegeGroup(mock.Anything).Return(nil).Times(len(util.BuiltinPrivilegeGroups))
+		meta.EXPECT().OperatePrivilegeGroup(mock.Anything, mock.Anything, mock.Anything).Return(nil).Times(len(util.BuiltinPrivilegeGroups))
+
+		Params.Save(Params.RbacConfig.Enabled.Key, "true")
+		Params.Save(Params.RbacConfig.ClusterReadWritePrivileges.Key, clusterReadWrite)
+
+		defer func() {
+			Params.Reset(Params.RbacConfig.Enabled.Key)
+			Params.Reset(Params.RbacConfig.ClusterReadWritePrivileges.Key)
+		}()
+
+		err := c.initBuiltinPrivilegeGroups()
+		assert.NoError(t, err)
+	})
 }
 
 func TestCore_BackupRBAC(t *testing.T) {