From 179e518c63b21ffe4b22a786a64d77f28b9bcc67 Mon Sep 17 00:00:00 2001 From: sthuang <167743503+shaoting-huang@users.noreply.github.com> Date: Tue, 3 Dec 2024 18:32:40 +0800 Subject: [PATCH] enhance: [GoSDK] add operate privilege & privilege group API (#38141) related issue: https://github.com/milvus-io/milvus/issues/37031 Signed-off-by: shaoting-huang --- client/entity/privilege_group.go | 23 ++++ client/go.mod | 2 +- client/go.sum | 4 +- client/mock_milvus_server_test.go | 55 ++++++++ client/rbac.go | 96 ++++++++++++++ client/rbac_options.go | 171 ++++++++++++++++++++++++ client/rbac_test.go | 209 ++++++++++++++++++++++++++++++ 7 files changed, 557 insertions(+), 3 deletions(-) create mode 100644 client/entity/privilege_group.go create mode 100644 client/rbac.go create mode 100644 client/rbac_options.go create mode 100644 client/rbac_test.go diff --git a/client/entity/privilege_group.go b/client/entity/privilege_group.go new file mode 100644 index 0000000000000..21162352e0a64 --- /dev/null +++ b/client/entity/privilege_group.go @@ -0,0 +1,23 @@ +// Licensed to the LF AI & Data foundation under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +package entity + +// PrivilegeGroup is the entity model for custom privilege group. +type PrivilegeGroup struct { + GroupName string + Privileges []string +} diff --git a/client/go.mod b/client/go.mod index 487253db7a8f4..d828b9b067f66 100644 --- a/client/go.mod +++ b/client/go.mod @@ -6,7 +6,7 @@ require ( github.com/blang/semver/v4 v4.0.0 github.com/cockroachdb/errors v1.9.1 github.com/grpc-ecosystem/go-grpc-middleware v1.3.0 - github.com/milvus-io/milvus-proto/go-api/v2 v2.4.17 + github.com/milvus-io/milvus-proto/go-api/v2 v2.4.18-0.20241120092224-a1c2ac2fd2c1 github.com/milvus-io/milvus/pkg v0.0.2-0.20240317152703-17b4938985f3 github.com/quasilyte/go-ruleguard/dsl v0.3.22 github.com/samber/lo v1.27.0 diff --git a/client/go.sum b/client/go.sum index f271e5f6ec2ec..24fc48b4979b8 100644 --- a/client/go.sum +++ b/client/go.sum @@ -400,8 +400,8 @@ github.com/matttproud/golang_protobuf_extensions v1.0.4/go.mod h1:BSXmuO+STAnVfr github.com/mediocregopher/radix/v3 v3.4.2/go.mod h1:8FL3F6UQRXHXIBSPUs5h0RybMF8i4n7wVopoX3x7Bv8= github.com/microcosm-cc/bluemonday v1.0.2/go.mod h1:iVP4YcDBq+n/5fb23BhYFvIMq/leAFZyRl6bYmGDlGc= github.com/miekg/dns v1.0.14/go.mod h1:W1PPwlIAgtquWBMBEV9nkV9Cazfe8ScdGz/Lj7v3Nrg= -github.com/milvus-io/milvus-proto/go-api/v2 v2.4.17 h1:ANkXdUKKpIPPQkw9pkV9ku9AEtSaPyua9XzdMTUxjCs= -github.com/milvus-io/milvus-proto/go-api/v2 v2.4.17/go.mod h1:/6UT4zZl6awVeXLeE7UGDWZvXj3IWkRsh3mqsn0DiAs= +github.com/milvus-io/milvus-proto/go-api/v2 v2.4.18-0.20241120092224-a1c2ac2fd2c1 h1:Xp4zOR85XFFtM7Eif945BeSmDf30hbdijbeNSuy92Bg= +github.com/milvus-io/milvus-proto/go-api/v2 v2.4.18-0.20241120092224-a1c2ac2fd2c1/go.mod h1:/6UT4zZl6awVeXLeE7UGDWZvXj3IWkRsh3mqsn0DiAs= github.com/milvus-io/milvus/pkg v0.0.2-0.20240317152703-17b4938985f3 h1:ZBpRWhBa7FTFxW4YYVv9AUESoW1Xyb3KNXTzTqfkZmw= github.com/milvus-io/milvus/pkg v0.0.2-0.20240317152703-17b4938985f3/go.mod h1:jQ2BUZny1COsgv1Qbcv8dmbppW+V9J/c4YQZNb3EOm8= github.com/mitchellh/cli v1.0.0/go.mod h1:hNIlj7HEI86fIcpObd7a0FcrxTWetlwJDGcceTlRvqc= diff --git a/client/mock_milvus_server_test.go b/client/mock_milvus_server_test.go index 3d8f94fde991f..e89a0bbb94724 100644 --- a/client/mock_milvus_server_test.go +++ b/client/mock_milvus_server_test.go @@ -4042,6 +4042,61 @@ func (_c *MilvusServiceServer_OperatePrivilegeGroup_Call) RunAndReturn(run func( return _c } +// OperatePrivilegeV2 provides a mock function with given fields: _a0, _a1 +func (_m *MilvusServiceServer) OperatePrivilegeV2(_a0 context.Context, _a1 *milvuspb.OperatePrivilegeV2Request) (*commonpb.Status, error) { + ret := _m.Called(_a0, _a1) + + var r0 *commonpb.Status + var r1 error + if rf, ok := ret.Get(0).(func(context.Context, *milvuspb.OperatePrivilegeV2Request) (*commonpb.Status, error)); ok { + return rf(_a0, _a1) + } + if rf, ok := ret.Get(0).(func(context.Context, *milvuspb.OperatePrivilegeV2Request) *commonpb.Status); ok { + r0 = rf(_a0, _a1) + } else { + if ret.Get(0) != nil { + r0 = ret.Get(0).(*commonpb.Status) + } + } + + if rf, ok := ret.Get(1).(func(context.Context, *milvuspb.OperatePrivilegeV2Request) error); ok { + r1 = rf(_a0, _a1) + } else { + r1 = ret.Error(1) + } + + return r0, r1 +} + +// MilvusServiceServer_OperatePrivilegeV2_Call is a *mock.Call that shadows Run/Return methods with type explicit version for method 'OperatePrivilegeV2' +type MilvusServiceServer_OperatePrivilegeV2_Call struct { + *mock.Call +} + +// OperatePrivilegeV2 is a helper method to define mock.On call +// - _a0 context.Context +// - _a1 *milvuspb.OperatePrivilegeV2Request +func (_e *MilvusServiceServer_Expecter) OperatePrivilegeV2(_a0 interface{}, _a1 interface{}) *MilvusServiceServer_OperatePrivilegeV2_Call { + return &MilvusServiceServer_OperatePrivilegeV2_Call{Call: _e.mock.On("OperatePrivilegeV2", _a0, _a1)} +} + +func (_c *MilvusServiceServer_OperatePrivilegeV2_Call) Run(run func(_a0 context.Context, _a1 *milvuspb.OperatePrivilegeV2Request)) *MilvusServiceServer_OperatePrivilegeV2_Call { + _c.Call.Run(func(args mock.Arguments) { + run(args[0].(context.Context), args[1].(*milvuspb.OperatePrivilegeV2Request)) + }) + return _c +} + +func (_c *MilvusServiceServer_OperatePrivilegeV2_Call) Return(_a0 *commonpb.Status, _a1 error) *MilvusServiceServer_OperatePrivilegeV2_Call { + _c.Call.Return(_a0, _a1) + return _c +} + +func (_c *MilvusServiceServer_OperatePrivilegeV2_Call) RunAndReturn(run func(context.Context, *milvuspb.OperatePrivilegeV2Request) (*commonpb.Status, error)) *MilvusServiceServer_OperatePrivilegeV2_Call { + _c.Call.Return(run) + return _c +} + // OperateUserRole provides a mock function with given fields: _a0, _a1 func (_m *MilvusServiceServer) OperateUserRole(_a0 context.Context, _a1 *milvuspb.OperateUserRoleRequest) (*commonpb.Status, error) { ret := _m.Called(_a0, _a1) diff --git a/client/rbac.go b/client/rbac.go new file mode 100644 index 0000000000000..c7c6323ab4654 --- /dev/null +++ b/client/rbac.go @@ -0,0 +1,96 @@ +// Licensed to the LF AI & Data foundation under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +package client + +import ( + "context" + + "github.com/samber/lo" + "google.golang.org/grpc" + + "github.com/milvus-io/milvus-proto/go-api/v2/milvuspb" + "github.com/milvus-io/milvus/client/v2/entity" + "github.com/milvus-io/milvus/pkg/util/merr" +) + +func (c *Client) GrantV2(ctx context.Context, option GrantV2Option, callOptions ...grpc.CallOption) error { + req := option.Request() + + return c.callService(func(milvusService milvuspb.MilvusServiceClient) error { + resp, err := milvusService.OperatePrivilegeV2(ctx, req, callOptions...) + return merr.CheckRPCCall(resp, err) + }) +} + +func (c *Client) RevokeV2(ctx context.Context, option RevokeV2Option, callOptions ...grpc.CallOption) error { + req := option.Request() + + return c.callService(func(milvusService milvuspb.MilvusServiceClient) error { + resp, err := milvusService.OperatePrivilegeV2(ctx, req, callOptions...) + return merr.CheckRPCCall(resp, err) + }) +} + +func (c *Client) CreatePrivilegeGroup(ctx context.Context, option CreatePrivilegeGroupOption, callOptions ...grpc.CallOption) error { + req := option.Request() + + return c.callService(func(milvusService milvuspb.MilvusServiceClient) error { + resp, err := milvusService.CreatePrivilegeGroup(ctx, req, callOptions...) + return merr.CheckRPCCall(resp, err) + }) +} + +func (c *Client) DropPrivilegeGroup(ctx context.Context, option DropPrivilegeGroupOption, callOptions ...grpc.CallOption) error { + req := option.Request() + + return c.callService(func(milvusService milvuspb.MilvusServiceClient) error { + resp, err := milvusService.DropPrivilegeGroup(ctx, req, callOptions...) + return merr.CheckRPCCall(resp, err) + }) +} + +func (c *Client) ListPrivilegeGroups(ctx context.Context, option ListPrivilegeGroupsOption, callOptions ...grpc.CallOption) ([]*entity.PrivilegeGroup, error) { + req := option.Request() + + var privilegeGroups []*entity.PrivilegeGroup + err := c.callService(func(milvusService milvuspb.MilvusServiceClient) error { + r, err := milvusService.ListPrivilegeGroups(ctx, req, callOptions...) + if err != nil { + return err + } + for _, pg := range r.PrivilegeGroups { + privileges := lo.Map(pg.Privileges, func(p *milvuspb.PrivilegeEntity, _ int) string { + return p.Name + }) + privilegeGroups = append(privilegeGroups, &entity.PrivilegeGroup{ + GroupName: pg.GroupName, + Privileges: privileges, + }) + } + return nil + }) + return privilegeGroups, err +} + +func (c *Client) OperatePrivilegeGroup(ctx context.Context, option OperatePrivilegeGroupOption, callOptions ...grpc.CallOption) error { + req := option.Request() + + return c.callService(func(milvusService milvuspb.MilvusServiceClient) error { + resp, err := milvusService.OperatePrivilegeGroup(ctx, req, callOptions...) + return merr.CheckRPCCall(resp, err) + }) +} diff --git a/client/rbac_options.go b/client/rbac_options.go new file mode 100644 index 0000000000000..8e54f44a13864 --- /dev/null +++ b/client/rbac_options.go @@ -0,0 +1,171 @@ +// Licensed to the LF AI & Data foundation under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +package client + +import ( + "github.com/milvus-io/milvus-proto/go-api/v2/milvuspb" +) + +// GrantV2Option is the interface builds OperatePrivilegeV2Request +type GrantV2Option interface { + Request() *milvuspb.OperatePrivilegeV2Request +} + +type grantV2Option struct { + roleName string + privilegeName string + dbName string + collectionName string +} + +func (opt *grantV2Option) Request() *milvuspb.OperatePrivilegeV2Request { + return &milvuspb.OperatePrivilegeV2Request{ + Role: &milvuspb.RoleEntity{Name: opt.roleName}, + Grantor: &milvuspb.GrantorEntity{ + Privilege: &milvuspb.PrivilegeEntity{Name: opt.privilegeName}, + }, + Type: milvuspb.OperatePrivilegeType_Grant, + DbName: opt.dbName, + CollectionName: opt.collectionName, + } +} + +func NewGrantV2Option(roleName, privilegeName, dbName, collectionName string) *grantV2Option { + return &grantV2Option{ + roleName: roleName, + privilegeName: privilegeName, + dbName: dbName, + collectionName: collectionName, + } +} + +// RevokeV2Option is the interface builds OperatePrivilegeV2Request +type RevokeV2Option interface { + Request() *milvuspb.OperatePrivilegeV2Request +} + +type revokeV2Option struct { + roleName string + privilegeName string + dbName string + collectionName string +} + +func (opt *revokeV2Option) Request() *milvuspb.OperatePrivilegeV2Request { + return &milvuspb.OperatePrivilegeV2Request{ + Role: &milvuspb.RoleEntity{Name: opt.roleName}, + Grantor: &milvuspb.GrantorEntity{ + Privilege: &milvuspb.PrivilegeEntity{Name: opt.privilegeName}, + }, + Type: milvuspb.OperatePrivilegeType_Revoke, + DbName: opt.dbName, + CollectionName: opt.collectionName, + } +} + +func NewRevokeV2Option(roleName, privilegeName, dbName, collectionName string) *revokeV2Option { + return &revokeV2Option{ + roleName: roleName, + privilegeName: privilegeName, + dbName: dbName, + collectionName: collectionName, + } +} + +// CreatePrivilegeGroupOption is the interface builds CreatePrivilegeGroupRequest +type CreatePrivilegeGroupOption interface { + Request() *milvuspb.CreatePrivilegeGroupRequest +} + +type createPrivilegeGroupOption struct { + groupName string +} + +func (opt *createPrivilegeGroupOption) Request() *milvuspb.CreatePrivilegeGroupRequest { + return &milvuspb.CreatePrivilegeGroupRequest{ + GroupName: opt.groupName, + } +} + +func NewCreatePrivilegeGroupOption(groupName string) *createPrivilegeGroupOption { + return &createPrivilegeGroupOption{ + groupName: groupName, + } +} + +// DropPrivilegeGroupOption is the interface builds DropPrivilegeGroupRequest +type DropPrivilegeGroupOption interface { + Request() *milvuspb.DropPrivilegeGroupRequest +} + +type dropPrivilegeGroupOption struct { + groupName string +} + +func (opt *dropPrivilegeGroupOption) Request() *milvuspb.DropPrivilegeGroupRequest { + return &milvuspb.DropPrivilegeGroupRequest{ + GroupName: opt.groupName, + } +} + +func NewDropPrivilegeGroupOption(groupName string) *dropPrivilegeGroupOption { + return &dropPrivilegeGroupOption{ + groupName: groupName, + } +} + +// ListPrivilegeGroupsOption is the interface builds ListPrivilegeGroupsRequest +type ListPrivilegeGroupsOption interface { + Request() *milvuspb.ListPrivilegeGroupsRequest +} + +type listPrivilegeGroupsOption struct{} + +func (opt *listPrivilegeGroupsOption) Request() *milvuspb.ListPrivilegeGroupsRequest { + return &milvuspb.ListPrivilegeGroupsRequest{} +} + +func NewListPrivilegeGroupsOption() *listPrivilegeGroupsOption { + return &listPrivilegeGroupsOption{} +} + +// OperatePrivilegeGroupOption is the interface builds OperatePrivilegeGroupRequest +type OperatePrivilegeGroupOption interface { + Request() *milvuspb.OperatePrivilegeGroupRequest +} + +type operatePrivilegeGroupOption struct { + groupName string + privileges []*milvuspb.PrivilegeEntity + operateType milvuspb.OperatePrivilegeGroupType +} + +func (opt *operatePrivilegeGroupOption) Request() *milvuspb.OperatePrivilegeGroupRequest { + return &milvuspb.OperatePrivilegeGroupRequest{ + GroupName: opt.groupName, + Privileges: opt.privileges, + Type: opt.operateType, + } +} + +func NewOperatePrivilegeGroupOption(groupName string, privileges []*milvuspb.PrivilegeEntity, operateType milvuspb.OperatePrivilegeGroupType) *operatePrivilegeGroupOption { + return &operatePrivilegeGroupOption{ + groupName: groupName, + privileges: privileges, + operateType: operateType, + } +} diff --git a/client/rbac_test.go b/client/rbac_test.go new file mode 100644 index 0000000000000..d0ab0fbee4bbe --- /dev/null +++ b/client/rbac_test.go @@ -0,0 +1,209 @@ +// Licensed to the LF AI & Data foundation under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +package client + +import ( + "context" + "fmt" + "testing" + + "github.com/stretchr/testify/mock" + "github.com/stretchr/testify/suite" + + "github.com/milvus-io/milvus-proto/go-api/v2/commonpb" + "github.com/milvus-io/milvus-proto/go-api/v2/milvuspb" + "github.com/milvus-io/milvus/pkg/util/merr" +) + +type PrivilgeGroupSuite struct { + MockSuiteBase +} + +func (s *PrivilgeGroupSuite) TestGrantV2() { + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + + roleName := fmt.Sprintf("test_role_%s", s.randString(6)) + privilegeName := "Insert" + dbName := fmt.Sprintf("test_db_%s", s.randString(6)) + collectionName := fmt.Sprintf("test_collection_%s", s.randString(6)) + + s.Run("success", func() { + s.mock.EXPECT().OperatePrivilegeV2(mock.Anything, mock.Anything).RunAndReturn(func(ctx context.Context, r *milvuspb.OperatePrivilegeV2Request) (*commonpb.Status, error) { + s.Equal(roleName, r.GetRole().GetName()) + s.Equal(privilegeName, r.GetGrantor().GetPrivilege().GetName()) + s.Equal(dbName, r.GetDbName()) + s.Equal(collectionName, r.GetCollectionName()) + return merr.Success(), nil + }).Once() + + err := s.client.GrantV2(ctx, NewGrantV2Option(roleName, privilegeName, dbName, collectionName)) + s.NoError(err) + }) + + s.Run("failure", func() { + s.mock.EXPECT().OperatePrivilegeV2(mock.Anything, mock.Anything).Return(nil, merr.WrapErrServiceInternal("mocked")).Once() + + err := s.client.GrantV2(ctx, NewGrantV2Option(roleName, privilegeName, dbName, collectionName)) + s.Error(err) + }) +} + +func (s *PrivilgeGroupSuite) TestRevokeV2() { + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + + roleName := fmt.Sprintf("test_role_%s", s.randString(6)) + privilegeName := "Insert" + dbName := fmt.Sprintf("test_db_%s", s.randString(6)) + collectionName := fmt.Sprintf("test_collection_%s", s.randString(6)) + + s.Run("success", func() { + s.mock.EXPECT().OperatePrivilegeV2(mock.Anything, mock.Anything).RunAndReturn(func(ctx context.Context, r *milvuspb.OperatePrivilegeV2Request) (*commonpb.Status, error) { + s.Equal(roleName, r.GetRole().GetName()) + s.Equal(privilegeName, r.GetGrantor().GetPrivilege().GetName()) + s.Equal(dbName, r.GetDbName()) + s.Equal(collectionName, r.GetCollectionName()) + return merr.Success(), nil + }).Once() + + err := s.client.RevokeV2(ctx, NewRevokeV2Option(roleName, privilegeName, dbName, collectionName)) + s.NoError(err) + }) + + s.Run("failure", func() { + s.mock.EXPECT().OperatePrivilegeV2(mock.Anything, mock.Anything).Return(nil, merr.WrapErrServiceInternal("mocked")).Once() + + err := s.client.RevokeV2(ctx, NewRevokeV2Option(roleName, privilegeName, dbName, collectionName)) + s.Error(err) + }) +} + +func (s *PrivilgeGroupSuite) TestCreatePrivilegeGroup() { + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + + groupName := fmt.Sprintf("test_pg_%s", s.randString(6)) + + s.Run("success", func() { + s.mock.EXPECT().CreatePrivilegeGroup(mock.Anything, mock.Anything).RunAndReturn(func(ctx context.Context, r *milvuspb.CreatePrivilegeGroupRequest) (*commonpb.Status, error) { + s.Equal(groupName, r.GetGroupName()) + return merr.Success(), nil + }).Once() + + err := s.client.CreatePrivilegeGroup(ctx, NewCreatePrivilegeGroupOption(groupName)) + s.NoError(err) + }) + + s.Run("failure", func() { + s.mock.EXPECT().CreatePrivilegeGroup(mock.Anything, mock.Anything).Return(nil, merr.WrapErrServiceInternal("mocked")).Once() + + err := s.client.CreatePrivilegeGroup(ctx, NewCreatePrivilegeGroupOption(groupName)) + s.Error(err) + }) +} + +func (s *PrivilgeGroupSuite) TestDropPrivilegeGroup() { + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + + groupName := fmt.Sprintf("test_pg_%s", s.randString(6)) + + s.Run("success", func() { + s.mock.EXPECT().DropPrivilegeGroup(mock.Anything, mock.Anything).RunAndReturn(func(ctx context.Context, r *milvuspb.DropPrivilegeGroupRequest) (*commonpb.Status, error) { + s.Equal(groupName, r.GetGroupName()) + return merr.Success(), nil + }).Once() + + err := s.client.DropPrivilegeGroup(ctx, NewDropPrivilegeGroupOption(groupName)) + s.NoError(err) + }) + + s.Run("failure", func() { + s.mock.EXPECT().DropPrivilegeGroup(mock.Anything, mock.Anything).Return(nil, merr.WrapErrServiceInternal("mocked")).Once() + + err := s.client.DropPrivilegeGroup(ctx, NewDropPrivilegeGroupOption(groupName)) + s.Error(err) + }) +} + +func (s *PrivilgeGroupSuite) TestListPrivilegeGroups() { + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + + s.Run("success", func() { + s.mock.EXPECT().ListPrivilegeGroups(mock.Anything, mock.Anything).RunAndReturn(func(ctx context.Context, r *milvuspb.ListPrivilegeGroupsRequest) (*milvuspb.ListPrivilegeGroupsResponse, error) { + return &milvuspb.ListPrivilegeGroupsResponse{ + PrivilegeGroups: []*milvuspb.PrivilegeGroupInfo{ + { + GroupName: "pg1", + Privileges: []*milvuspb.PrivilegeEntity{{Name: "Insert"}, {Name: "Query"}}, + }, + { + GroupName: "pg2", + Privileges: []*milvuspb.PrivilegeEntity{{Name: "Delete"}, {Name: "Query"}}, + }, + }, + }, nil + }).Once() + + pgs, err := s.client.ListPrivilegeGroups(ctx, NewListPrivilegeGroupsOption()) + s.NoError(err) + s.Equal(2, len(pgs)) + s.Equal("pg1", pgs[0].GroupName) + s.Equal([]string{"Insert", "Query"}, pgs[0].Privileges) + s.Equal("pg2", pgs[1].GroupName) + s.Equal([]string{"Delete", "Query"}, pgs[1].Privileges) + }) + + s.Run("failure", func() { + s.mock.EXPECT().ListPrivilegeGroups(mock.Anything, mock.Anything).Return(nil, merr.WrapErrServiceInternal("mocked")).Once() + + _, err := s.client.ListPrivilegeGroups(ctx, NewListPrivilegeGroupsOption()) + s.Error(err) + }) +} + +func (s *PrivilgeGroupSuite) TestOperatePrivilegeGroup() { + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + + groupName := fmt.Sprintf("test_pg_%s", s.randString(6)) + privileges := []*milvuspb.PrivilegeEntity{{Name: "Insert"}, {Name: "Query"}} + operateType := milvuspb.OperatePrivilegeGroupType_AddPrivilegesToGroup + + s.Run("success", func() { + s.mock.EXPECT().OperatePrivilegeGroup(mock.Anything, mock.Anything).RunAndReturn(func(ctx context.Context, r *milvuspb.OperatePrivilegeGroupRequest) (*commonpb.Status, error) { + s.Equal(groupName, r.GetGroupName()) + return merr.Success(), nil + }).Once() + + err := s.client.OperatePrivilegeGroup(ctx, NewOperatePrivilegeGroupOption(groupName, privileges, operateType)) + s.NoError(err) + }) + + s.Run("failure", func() { + s.mock.EXPECT().OperatePrivilegeGroup(mock.Anything, mock.Anything).Return(nil, merr.WrapErrServiceInternal("mocked")).Once() + + err := s.client.OperatePrivilegeGroup(ctx, NewOperatePrivilegeGroupOption(groupName, privileges, operateType)) + s.Error(err) + }) +} + +func TestPrivilegeGroup(t *testing.T) { + suite.Run(t, new(PrivilgeGroupSuite)) +}