github.com/milvus-io/milvus-sdk-go/v2@v2.4.1/client/rbac_test.go (about)

     1  package client
     2  
     3  import (
     4  	"context"
     5  	"testing"
     6  
     7  	"github.com/cockroachdb/errors"
     8  
     9  	"github.com/milvus-io/milvus-sdk-go/v2/entity"
    10  	"github.com/stretchr/testify/suite"
    11  
    12  	"github.com/stretchr/testify/mock"
    13  
    14  	"github.com/milvus-io/milvus-proto/go-api/v2/commonpb"
    15  	"github.com/milvus-io/milvus-proto/go-api/v2/milvuspb"
    16  )
    17  
    18  type RBACSuite struct {
    19  	MockSuiteBase
    20  }
    21  
    22  func (s *RBACSuite) TestCreateRole() {
    23  	ctx := context.Background()
    24  	roleName := "testRole"
    25  
    26  	s.Run("normal run", func() {
    27  		ctx, cancel := context.WithCancel(ctx)
    28  		defer cancel()
    29  		defer s.resetMock()
    30  		s.mock.EXPECT().CreateRole(mock.Anything, mock.Anything).Run(func(ctx context.Context, req *milvuspb.CreateRoleRequest) {
    31  			s.Equal(roleName, req.GetEntity().GetName())
    32  		}).Return(&commonpb.Status{
    33  			ErrorCode: commonpb.ErrorCode_Success,
    34  		}, nil)
    35  		err := s.client.CreateRole(ctx, roleName)
    36  
    37  		s.NoError(err)
    38  	})
    39  
    40  	s.Run("rpc error", func() {
    41  		ctx, cancel := context.WithCancel(ctx)
    42  		defer cancel()
    43  		defer s.resetMock()
    44  		s.mock.EXPECT().CreateRole(mock.Anything, mock.Anything).Return(nil, errors.New("mock error"))
    45  		err := s.client.CreateRole(ctx, roleName)
    46  
    47  		s.Error(err)
    48  	})
    49  
    50  	s.Run("status error", func() {
    51  		ctx, cancel := context.WithCancel(ctx)
    52  		defer cancel()
    53  		defer s.resetMock()
    54  		s.mock.EXPECT().CreateRole(mock.Anything, mock.Anything).Return(&commonpb.Status{
    55  			ErrorCode: commonpb.ErrorCode_UnexpectedError,
    56  		}, nil)
    57  		err := s.client.CreateRole(ctx, roleName)
    58  
    59  		s.Error(err)
    60  	})
    61  
    62  	s.Run("service not ready", func() {
    63  		ctx, cancel := context.WithCancel(ctx)
    64  		defer cancel()
    65  
    66  		c := &GrpcClient{}
    67  		err := c.CreateRole(ctx, roleName)
    68  		s.Error(err)
    69  		s.ErrorIs(err, ErrClientNotReady)
    70  	})
    71  }
    72  
    73  func (s *RBACSuite) TestDropRole() {
    74  	ctx := context.Background()
    75  	roleName := "testRole"
    76  
    77  	s.Run("normal run", func() {
    78  		ctx, cancel := context.WithCancel(ctx)
    79  		defer cancel()
    80  		defer s.resetMock()
    81  		s.mock.EXPECT().DropRole(mock.Anything, mock.Anything).Run(func(ctx context.Context, req *milvuspb.DropRoleRequest) {
    82  			s.Equal(roleName, req.GetRoleName())
    83  		}).Return(&commonpb.Status{
    84  			ErrorCode: commonpb.ErrorCode_Success,
    85  		}, nil)
    86  		err := s.client.DropRole(ctx, roleName)
    87  
    88  		s.NoError(err)
    89  	})
    90  
    91  	s.Run("rpc error", func() {
    92  		ctx, cancel := context.WithCancel(ctx)
    93  		defer cancel()
    94  		defer s.resetMock()
    95  		s.mock.EXPECT().DropRole(mock.Anything, mock.Anything).Return(nil, errors.New("mock error"))
    96  		err := s.client.DropRole(ctx, roleName)
    97  
    98  		s.Error(err)
    99  	})
   100  
   101  	s.Run("status error", func() {
   102  		ctx, cancel := context.WithCancel(ctx)
   103  		defer cancel()
   104  		defer s.resetMock()
   105  		s.mock.EXPECT().DropRole(mock.Anything, mock.Anything).Return(&commonpb.Status{
   106  			ErrorCode: commonpb.ErrorCode_UnexpectedError,
   107  		}, nil)
   108  		err := s.client.DropRole(ctx, roleName)
   109  
   110  		s.Error(err)
   111  	})
   112  
   113  	s.Run("service not ready", func() {
   114  		ctx, cancel := context.WithCancel(ctx)
   115  		defer cancel()
   116  
   117  		c := &GrpcClient{}
   118  		err := c.DropRole(ctx, roleName)
   119  		s.Error(err)
   120  		s.ErrorIs(err, ErrClientNotReady)
   121  	})
   122  }
   123  
   124  func (s *RBACSuite) TestAddUserRole() {
   125  	ctx := context.Background()
   126  	username := "testUser"
   127  	roleName := "testRole"
   128  
   129  	s.Run("normal run", func() {
   130  		ctx, cancel := context.WithCancel(ctx)
   131  		defer cancel()
   132  		defer s.resetMock()
   133  		s.mock.EXPECT().OperateUserRole(mock.Anything, mock.Anything).Run(func(ctx context.Context, req *milvuspb.OperateUserRoleRequest) {
   134  			s.Equal(milvuspb.OperateUserRoleType_AddUserToRole, req.GetType())
   135  			s.Equal(username, req.GetUsername())
   136  			s.Equal(roleName, req.GetRoleName())
   137  		}).Return(&commonpb.Status{
   138  			ErrorCode: commonpb.ErrorCode_Success,
   139  		}, nil)
   140  		err := s.client.AddUserRole(ctx, username, roleName)
   141  
   142  		s.NoError(err)
   143  	})
   144  
   145  	s.Run("rpc error", func() {
   146  		ctx, cancel := context.WithCancel(ctx)
   147  		defer cancel()
   148  		defer s.resetMock()
   149  		s.mock.EXPECT().OperateUserRole(mock.Anything, mock.Anything).Return(nil, errors.New("mock error"))
   150  		err := s.client.AddUserRole(ctx, username, roleName)
   151  
   152  		s.Error(err)
   153  	})
   154  
   155  	s.Run("status error", func() {
   156  		ctx, cancel := context.WithCancel(ctx)
   157  		defer cancel()
   158  		defer s.resetMock()
   159  		s.mock.EXPECT().OperateUserRole(mock.Anything, mock.Anything).Return(&commonpb.Status{
   160  			ErrorCode: commonpb.ErrorCode_UnexpectedError,
   161  		}, nil)
   162  		err := s.client.AddUserRole(ctx, username, roleName)
   163  
   164  		s.Error(err)
   165  	})
   166  
   167  	s.Run("service not ready", func() {
   168  		ctx, cancel := context.WithCancel(ctx)
   169  		defer cancel()
   170  
   171  		c := &GrpcClient{}
   172  		err := c.AddUserRole(ctx, username, roleName)
   173  		s.Error(err)
   174  		s.ErrorIs(err, ErrClientNotReady)
   175  	})
   176  }
   177  
   178  func (s *RBACSuite) TestRemoveUserRole() {
   179  	ctx := context.Background()
   180  	username := "testUser"
   181  	roleName := "testRole"
   182  
   183  	s.Run("normal run", func() {
   184  		ctx, cancel := context.WithCancel(ctx)
   185  		defer cancel()
   186  		defer s.resetMock()
   187  		s.mock.EXPECT().OperateUserRole(mock.Anything, mock.Anything).Run(func(ctx context.Context, req *milvuspb.OperateUserRoleRequest) {
   188  			s.Equal(milvuspb.OperateUserRoleType_RemoveUserFromRole, req.GetType())
   189  			s.Equal(username, req.GetUsername())
   190  			s.Equal(roleName, req.GetRoleName())
   191  		}).Return(&commonpb.Status{
   192  			ErrorCode: commonpb.ErrorCode_Success,
   193  		}, nil)
   194  		err := s.client.RemoveUserRole(ctx, username, roleName)
   195  
   196  		s.NoError(err)
   197  	})
   198  
   199  	s.Run("rpc error", func() {
   200  		ctx, cancel := context.WithCancel(ctx)
   201  		defer cancel()
   202  		defer s.resetMock()
   203  		s.mock.EXPECT().OperateUserRole(mock.Anything, mock.Anything).Return(nil, errors.New("mock error"))
   204  		err := s.client.RemoveUserRole(ctx, username, roleName)
   205  
   206  		s.Error(err)
   207  	})
   208  
   209  	s.Run("status error", func() {
   210  		ctx, cancel := context.WithCancel(ctx)
   211  		defer cancel()
   212  		defer s.resetMock()
   213  		s.mock.EXPECT().OperateUserRole(mock.Anything, mock.Anything).Return(&commonpb.Status{
   214  			ErrorCode: commonpb.ErrorCode_UnexpectedError,
   215  		}, nil)
   216  		err := s.client.RemoveUserRole(ctx, username, roleName)
   217  
   218  		s.Error(err)
   219  	})
   220  
   221  	s.Run("service not ready", func() {
   222  		ctx, cancel := context.WithCancel(ctx)
   223  		defer cancel()
   224  
   225  		c := &GrpcClient{}
   226  		err := c.RemoveUserRole(ctx, username, roleName)
   227  		s.Error(err)
   228  		s.ErrorIs(err, ErrClientNotReady)
   229  	})
   230  }
   231  
   232  func (s *RBACSuite) TestListRoles() {
   233  	ctx := context.Background()
   234  	roleName := "testRole"
   235  
   236  	s.Run("normal run", func() {
   237  		ctx, cancel := context.WithCancel(ctx)
   238  		defer cancel()
   239  		defer s.resetMock()
   240  		s.mock.EXPECT().SelectRole(mock.Anything, mock.Anything).Run(func(ctx context.Context, req *milvuspb.SelectRoleRequest) {
   241  			s.False(req.GetIncludeUserInfo())
   242  		}).Return(&milvuspb.SelectRoleResponse{
   243  			Status: &commonpb.Status{ErrorCode: commonpb.ErrorCode_Success},
   244  			Results: []*milvuspb.RoleResult{
   245  				{
   246  					Role: &milvuspb.RoleEntity{
   247  						Name: roleName,
   248  					},
   249  				},
   250  			},
   251  		}, nil)
   252  		roles, err := s.client.ListRoles(ctx)
   253  
   254  		s.NoError(err)
   255  		s.ElementsMatch([]entity.Role{{Name: roleName}}, roles)
   256  	})
   257  
   258  	s.Run("rpc error", func() {
   259  		ctx, cancel := context.WithCancel(ctx)
   260  		defer cancel()
   261  		defer s.resetMock()
   262  		s.mock.EXPECT().SelectRole(mock.Anything, mock.Anything).Return(nil, errors.New("mock error"))
   263  
   264  		_, err := s.client.ListRoles(ctx)
   265  		s.Error(err)
   266  	})
   267  
   268  	s.Run("status error", func() {
   269  		ctx, cancel := context.WithCancel(ctx)
   270  		defer cancel()
   271  		defer s.resetMock()
   272  		s.mock.EXPECT().SelectRole(mock.Anything, mock.Anything).Return(&milvuspb.SelectRoleResponse{
   273  			Status: &commonpb.Status{ErrorCode: commonpb.ErrorCode_UnexpectedError},
   274  		}, nil)
   275  
   276  		_, err := s.client.ListRoles(ctx)
   277  		s.Error(err)
   278  	})
   279  
   280  	s.Run("service not ready", func() {
   281  		ctx, cancel := context.WithCancel(ctx)
   282  		defer cancel()
   283  
   284  		c := &GrpcClient{}
   285  		_, err := c.ListRoles(ctx)
   286  		s.Error(err)
   287  		s.ErrorIs(err, ErrClientNotReady)
   288  	})
   289  }
   290  
   291  func (s *RBACSuite) TestListUser() {
   292  	ctx := context.Background()
   293  	userName := "testUser"
   294  
   295  	s.Run("normal run", func() {
   296  		ctx, cancel := context.WithCancel(ctx)
   297  		defer cancel()
   298  		defer s.resetMock()
   299  		s.mock.EXPECT().SelectUser(mock.Anything, mock.Anything).Run(func(ctx context.Context, req *milvuspb.SelectUserRequest) {
   300  			s.False(req.GetIncludeRoleInfo())
   301  		}).Return(&milvuspb.SelectUserResponse{
   302  			Status: &commonpb.Status{ErrorCode: commonpb.ErrorCode_Success},
   303  			Results: []*milvuspb.UserResult{
   304  				{
   305  					User: &milvuspb.UserEntity{
   306  						Name: userName,
   307  					},
   308  				},
   309  			},
   310  		}, nil)
   311  		users, err := s.client.ListUsers(ctx)
   312  
   313  		s.NoError(err)
   314  		s.ElementsMatch([]entity.User{{Name: userName}}, users)
   315  	})
   316  
   317  	s.Run("rpc error", func() {
   318  		ctx, cancel := context.WithCancel(ctx)
   319  		defer cancel()
   320  		defer s.resetMock()
   321  		s.mock.EXPECT().SelectUser(mock.Anything, mock.Anything).Return(nil, errors.New("mock error"))
   322  
   323  		_, err := s.client.ListUsers(ctx)
   324  		s.Error(err)
   325  	})
   326  
   327  	s.Run("status error", func() {
   328  		ctx, cancel := context.WithCancel(ctx)
   329  		defer cancel()
   330  		defer s.resetMock()
   331  		s.mock.EXPECT().SelectUser(mock.Anything, mock.Anything).Return(&milvuspb.SelectUserResponse{
   332  			Status: &commonpb.Status{ErrorCode: commonpb.ErrorCode_UnexpectedError},
   333  		}, nil)
   334  
   335  		_, err := s.client.ListUsers(ctx)
   336  		s.Error(err)
   337  	})
   338  
   339  	s.Run("service not ready", func() {
   340  		ctx, cancel := context.WithCancel(ctx)
   341  		defer cancel()
   342  
   343  		c := &GrpcClient{}
   344  		_, err := c.ListUsers(ctx)
   345  		s.Error(err)
   346  		s.ErrorIs(err, ErrClientNotReady)
   347  	})
   348  }
   349  
   350  func (s *RBACSuite) TestDescribeUser() {
   351  	ctx := context.Background()
   352  	userName := "testUser"
   353  
   354  	s.Run("normal run", func() {
   355  		ctx, cancel := context.WithCancel(ctx)
   356  		defer cancel()
   357  		defer s.resetMock()
   358  		s.mock.EXPECT().SelectUser(mock.Anything, mock.Anything).Run(func(ctx context.Context, req *milvuspb.SelectUserRequest) {
   359  			s.True(req.GetIncludeRoleInfo())
   360  			s.Equal(req.GetUser().GetName(), userName)
   361  		}).Return(&milvuspb.SelectUserResponse{
   362  			Status: &commonpb.Status{ErrorCode: commonpb.ErrorCode_Success},
   363  			Results: []*milvuspb.UserResult{
   364  				{
   365  					User: &milvuspb.UserEntity{
   366  						Name: userName,
   367  					},
   368  					Roles: []*milvuspb.RoleEntity{
   369  						{Name: "role1"},
   370  						{Name: "role2"},
   371  					},
   372  				},
   373  			},
   374  		}, nil)
   375  
   376  		userDesc, err := s.client.DescribeUser(ctx, userName)
   377  
   378  		s.NoError(err)
   379  		s.Equal(userDesc.Name, userName)
   380  		s.ElementsMatch(userDesc.Roles, []string{"role1", "role2"})
   381  	})
   382  
   383  	s.Run("rpc error", func() {
   384  		ctx, cancel := context.WithCancel(ctx)
   385  		defer cancel()
   386  		defer s.resetMock()
   387  		s.mock.EXPECT().SelectUser(mock.Anything, mock.Anything).Return(nil, errors.New("mock error"))
   388  
   389  		_, err := s.client.DescribeUser(ctx, userName)
   390  		s.Error(err)
   391  	})
   392  
   393  	s.Run("status error", func() {
   394  		ctx, cancel := context.WithCancel(ctx)
   395  		defer cancel()
   396  		defer s.resetMock()
   397  		s.mock.EXPECT().SelectUser(mock.Anything, mock.Anything).Return(&milvuspb.SelectUserResponse{
   398  			Status: &commonpb.Status{ErrorCode: commonpb.ErrorCode_UnexpectedError},
   399  		}, nil)
   400  
   401  		_, err := s.client.DescribeUser(ctx, userName)
   402  		s.Error(err)
   403  	})
   404  
   405  	s.Run("service not ready", func() {
   406  		ctx, cancel := context.WithCancel(ctx)
   407  		defer cancel()
   408  
   409  		c := &GrpcClient{}
   410  		_, err := c.DescribeUser(ctx, userName)
   411  		s.Error(err)
   412  		s.ErrorIs(err, ErrClientNotReady)
   413  	})
   414  }
   415  
   416  func (s *RBACSuite) TestDescribeUsers() {
   417  	ctx := context.Background()
   418  
   419  	s.Run("normal run", func() {
   420  		ctx, cancel := context.WithCancel(ctx)
   421  		defer cancel()
   422  		defer s.resetMock()
   423  
   424  		mockResults := []*milvuspb.UserResult{
   425  			{
   426  				User: &milvuspb.UserEntity{
   427  					Name: "user1",
   428  				},
   429  				Roles: []*milvuspb.RoleEntity{
   430  					{Name: "role1"},
   431  					{Name: "role2"},
   432  				},
   433  			},
   434  			{
   435  				User: &milvuspb.UserEntity{
   436  					Name: "user2",
   437  				},
   438  				Roles: []*milvuspb.RoleEntity{
   439  					{Name: "role3"},
   440  				},
   441  			},
   442  		}
   443  
   444  		s.mock.EXPECT().SelectUser(mock.Anything, mock.Anything).Run(func(ctx context.Context, req *milvuspb.SelectUserRequest) {
   445  			s.True(req.GetIncludeRoleInfo())
   446  		}).Return(&milvuspb.SelectUserResponse{
   447  			Status:  &commonpb.Status{ErrorCode: commonpb.ErrorCode_Success},
   448  			Results: mockResults,
   449  		}, nil)
   450  
   451  		userDescs, err := s.client.DescribeUsers(ctx)
   452  
   453  		s.NoError(err)
   454  		s.Len(userDescs, 2)
   455  
   456  		expectedDescs := []entity.UserDescription{
   457  			{
   458  				Name:  "user1",
   459  				Roles: []string{"role1", "role2"},
   460  			},
   461  			{
   462  				Name:  "user2",
   463  				Roles: []string{"role3"},
   464  			},
   465  		}
   466  
   467  		for i, desc := range userDescs {
   468  			s.Equal(expectedDescs[i].Name, desc.Name)
   469  			s.ElementsMatch(expectedDescs[i].Roles, desc.Roles)
   470  		}
   471  	})
   472  
   473  	s.Run("rpc error", func() {
   474  		ctx, cancel := context.WithCancel(ctx)
   475  		defer cancel()
   476  		defer s.resetMock()
   477  
   478  		s.mock.EXPECT().SelectUser(mock.Anything, mock.Anything).Return(nil, errors.New("mock error"))
   479  
   480  		_, err := s.client.DescribeUsers(ctx)
   481  		s.Error(err)
   482  	})
   483  
   484  	s.Run("status error", func() {
   485  		ctx, cancel := context.WithCancel(ctx)
   486  		defer cancel()
   487  		defer s.resetMock()
   488  
   489  		s.mock.EXPECT().SelectUser(mock.Anything, mock.Anything).Return(&milvuspb.SelectUserResponse{
   490  			Status: &commonpb.Status{ErrorCode: commonpb.ErrorCode_UnexpectedError},
   491  		}, nil)
   492  
   493  		_, err := s.client.DescribeUsers(ctx)
   494  		s.Error(err)
   495  	})
   496  
   497  	s.Run("service not ready", func() {
   498  		ctx, cancel := context.WithCancel(ctx)
   499  		defer cancel()
   500  
   501  		c := &GrpcClient{}
   502  		_, err := c.DescribeUsers(ctx)
   503  		s.Error(err)
   504  		s.ErrorIs(err, ErrClientNotReady)
   505  	})
   506  }
   507  
   508  func (s *RBACSuite) TestListGrant() {
   509  	ctx := context.Background()
   510  	roleName := "testRole"
   511  	object := "testObject"
   512  	objectName := "testObjectName"
   513  	dbName := "testDB"
   514  
   515  	s.Run("normal run", func() {
   516  		ctx, cancel := context.WithCancel(ctx)
   517  		defer cancel()
   518  		defer s.resetMock()
   519  
   520  		mockResults := []*milvuspb.GrantEntity{
   521  			{
   522  				Object: &milvuspb.ObjectEntity{
   523  					Name: object,
   524  				},
   525  				ObjectName: objectName,
   526  				Role: &milvuspb.RoleEntity{
   527  					Name: roleName,
   528  				},
   529  				Grantor: &milvuspb.GrantorEntity{
   530  					User: &milvuspb.UserEntity{
   531  						Name: "grantorUser",
   532  					},
   533  					Privilege: &milvuspb.PrivilegeEntity{
   534  						Name: "testPrivilege",
   535  					},
   536  				},
   537  				DbName: dbName,
   538  			},
   539  		}
   540  
   541  		s.mock.EXPECT().SelectGrant(mock.Anything, mock.Anything).Run(func(ctx context.Context, req *milvuspb.SelectGrantRequest) {
   542  			s.Equal(req.GetEntity().GetRole().GetName(), roleName)
   543  			s.Equal(req.GetEntity().GetObject().GetName(), object)
   544  			s.Equal(req.GetEntity().GetObjectName(), objectName)
   545  			s.Equal(req.GetEntity().GetDbName(), dbName)
   546  		}).Return(&milvuspb.SelectGrantResponse{
   547  			Status:   &commonpb.Status{ErrorCode: commonpb.ErrorCode_Success},
   548  			Entities: mockResults,
   549  		}, nil)
   550  
   551  		grants, err := s.client.ListGrant(ctx, roleName, object, objectName, dbName)
   552  
   553  		s.NoError(err)
   554  		s.Len(grants, 1)
   555  
   556  		expectedGrant := entity.RoleGrants{
   557  			Object:        object,
   558  			ObjectName:    objectName,
   559  			RoleName:      roleName,
   560  			GrantorName:   "grantorUser",
   561  			PrivilegeName: "testPrivilege",
   562  			DbName:        dbName,
   563  		}
   564  
   565  		s.Equal(expectedGrant, grants[0])
   566  	})
   567  }
   568  
   569  func (s *RBACSuite) TestListGrants() {
   570  	ctx := context.Background()
   571  	roleName := "testRole"
   572  	dbName := "testDB"
   573  
   574  	s.Run("normal run", func() {
   575  		ctx, cancel := context.WithCancel(ctx)
   576  		defer cancel()
   577  		defer s.resetMock()
   578  
   579  		mockResults := []*milvuspb.GrantEntity{
   580  			{
   581  				Object: &milvuspb.ObjectEntity{
   582  					Name: "testObject",
   583  				},
   584  				ObjectName: "testObjectName",
   585  				Role: &milvuspb.RoleEntity{
   586  					Name: roleName,
   587  				},
   588  				Grantor: &milvuspb.GrantorEntity{
   589  					User: &milvuspb.UserEntity{
   590  						Name: "grantorUser",
   591  					},
   592  					Privilege: &milvuspb.PrivilegeEntity{
   593  						Name: "testPrivilege",
   594  					},
   595  				},
   596  				DbName: dbName,
   597  			},
   598  		}
   599  
   600  		s.mock.EXPECT().SelectGrant(mock.Anything, mock.Anything).Run(func(ctx context.Context, req *milvuspb.SelectGrantRequest) {
   601  			s.Equal(req.GetEntity().GetRole().GetName(), roleName)
   602  			s.Equal(req.GetEntity().GetDbName(), dbName)
   603  		}).Return(&milvuspb.SelectGrantResponse{
   604  			Status:   &commonpb.Status{ErrorCode: commonpb.ErrorCode_Success},
   605  			Entities: mockResults,
   606  		}, nil)
   607  
   608  		grants, err := s.client.ListGrants(ctx, roleName, dbName)
   609  
   610  		s.NoError(err)
   611  		s.Len(grants, 1)
   612  
   613  		expectedGrant := entity.RoleGrants{
   614  			Object:        "testObject",
   615  			ObjectName:    "testObjectName",
   616  			RoleName:      roleName,
   617  			GrantorName:   "grantorUser",
   618  			PrivilegeName: "testPrivilege",
   619  			DbName:        dbName,
   620  		}
   621  
   622  		s.Equal(expectedGrant, grants[0])
   623  	})
   624  }
   625  
   626  func (s *RBACSuite) TestGrant() {
   627  	ctx := context.Background()
   628  
   629  	roleName := "testRole"
   630  	objectName := testCollectionName
   631  	objectType := entity.PriviledegeObjectTypeCollection
   632  
   633  	s.Run("normal run", func() {
   634  		ctx, cancel := context.WithCancel(ctx)
   635  		defer cancel()
   636  		defer s.resetMock()
   637  		s.mock.EXPECT().OperatePrivilege(mock.Anything, mock.Anything).Run(func(ctx context.Context, req *milvuspb.OperatePrivilegeRequest) {
   638  			s.Equal(roleName, req.GetEntity().GetRole().GetName())
   639  			s.Equal(objectName, req.GetEntity().GetObjectName())
   640  			s.Equal(commonpb.ObjectType_name[int32(objectType)], req.GetEntity().GetObject().GetName())
   641  			s.Equal(milvuspb.OperatePrivilegeType_Grant, req.GetType())
   642  		}).Return(&commonpb.Status{ErrorCode: commonpb.ErrorCode_Success}, nil)
   643  
   644  		err := s.client.Grant(ctx, roleName, objectType, objectName)
   645  
   646  		s.NoError(err)
   647  	})
   648  
   649  	s.Run("rpc error", func() {
   650  		ctx, cancel := context.WithCancel(ctx)
   651  		defer cancel()
   652  		defer s.resetMock()
   653  		s.mock.EXPECT().OperatePrivilege(mock.Anything, mock.Anything).Return(nil, errors.New("mock error"))
   654  
   655  		err := s.client.Grant(ctx, roleName, objectType, objectName)
   656  		s.Error(err)
   657  	})
   658  
   659  	s.Run("status error", func() {
   660  		ctx, cancel := context.WithCancel(ctx)
   661  		defer cancel()
   662  		defer s.resetMock()
   663  		s.mock.EXPECT().OperatePrivilege(mock.Anything, mock.Anything).Return(&commonpb.Status{ErrorCode: commonpb.ErrorCode_UnexpectedError}, nil)
   664  
   665  		err := s.client.Grant(ctx, roleName, objectType, objectName)
   666  		s.Error(err)
   667  	})
   668  
   669  	s.Run("service not ready", func() {
   670  		ctx, cancel := context.WithCancel(ctx)
   671  		defer cancel()
   672  
   673  		c := &GrpcClient{}
   674  		err := c.Grant(ctx, roleName, objectType, objectName)
   675  		s.Error(err)
   676  		s.ErrorIs(err, ErrClientNotReady)
   677  	})
   678  }
   679  
   680  func (s *RBACSuite) TestRevoke() {
   681  	ctx := context.Background()
   682  
   683  	roleName := "testRole"
   684  	objectName := testCollectionName
   685  	objectType := entity.PriviledegeObjectTypeCollection
   686  
   687  	s.Run("normal run", func() {
   688  		ctx, cancel := context.WithCancel(ctx)
   689  		defer cancel()
   690  		defer s.resetMock()
   691  		s.mock.EXPECT().OperatePrivilege(mock.Anything, mock.Anything).Run(func(ctx context.Context, req *milvuspb.OperatePrivilegeRequest) {
   692  			s.Equal(roleName, req.GetEntity().GetRole().GetName())
   693  			s.Equal(objectName, req.GetEntity().GetObjectName())
   694  			s.Equal(commonpb.ObjectType_name[int32(objectType)], req.GetEntity().GetObject().GetName())
   695  			s.Equal(milvuspb.OperatePrivilegeType_Revoke, req.GetType())
   696  		}).Return(&commonpb.Status{ErrorCode: commonpb.ErrorCode_Success}, nil)
   697  
   698  		err := s.client.Revoke(ctx, roleName, objectType, objectName)
   699  
   700  		s.NoError(err)
   701  	})
   702  
   703  	s.Run("rpc error", func() {
   704  		ctx, cancel := context.WithCancel(ctx)
   705  		defer cancel()
   706  		defer s.resetMock()
   707  		s.mock.EXPECT().OperatePrivilege(mock.Anything, mock.Anything).Return(nil, errors.New("mock error"))
   708  
   709  		err := s.client.Revoke(ctx, roleName, objectType, objectName)
   710  		s.Error(err)
   711  	})
   712  
   713  	s.Run("status error", func() {
   714  		ctx, cancel := context.WithCancel(ctx)
   715  		defer cancel()
   716  		defer s.resetMock()
   717  		s.mock.EXPECT().OperatePrivilege(mock.Anything, mock.Anything).Return(&commonpb.Status{ErrorCode: commonpb.ErrorCode_UnexpectedError}, nil)
   718  
   719  		err := s.client.Revoke(ctx, roleName, objectType, objectName)
   720  		s.Error(err)
   721  	})
   722  
   723  	s.Run("service not ready", func() {
   724  		ctx, cancel := context.WithCancel(ctx)
   725  		defer cancel()
   726  		defer s.resetMock()
   727  
   728  		c := &GrpcClient{}
   729  		err := c.Revoke(ctx, roleName, objectType, objectName)
   730  		s.Error(err)
   731  		s.ErrorIs(err, ErrClientNotReady)
   732  	})
   733  }
   734  
   735  func TestRBACSuite(t *testing.T) {
   736  	suite.Run(t, new(RBACSuite))
   737  }