database_test.go 2.7 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293
  1. package client
  2. import (
  3. "context"
  4. "fmt"
  5. "testing"
  6. mock "github.com/stretchr/testify/mock"
  7. "github.com/stretchr/testify/suite"
  8. "github.com/milvus-io/milvus-proto/go-api/v2/commonpb"
  9. "github.com/milvus-io/milvus-proto/go-api/v2/milvuspb"
  10. "github.com/milvus-io/milvus/pkg/util/merr"
  11. )
  12. type DatabaseSuite struct {
  13. MockSuiteBase
  14. }
  15. func (s *DatabaseSuite) TestListDatabases() {
  16. ctx, cancel := context.WithCancel(context.Background())
  17. defer cancel()
  18. s.Run("success", func() {
  19. s.mock.EXPECT().ListDatabases(mock.Anything, mock.Anything).Return(&milvuspb.ListDatabasesResponse{
  20. Status: merr.Success(),
  21. DbNames: []string{"default", "db1"},
  22. }, nil).Once()
  23. names, err := s.client.ListDatabase(ctx, NewListDatabaseOption())
  24. s.NoError(err)
  25. s.ElementsMatch([]string{"default", "db1"}, names)
  26. })
  27. s.Run("failure", func() {
  28. s.mock.EXPECT().ListDatabases(mock.Anything, mock.Anything).Return(nil, merr.WrapErrServiceInternal("mocked")).Once()
  29. _, err := s.client.ListDatabase(ctx, NewListDatabaseOption())
  30. s.Error(err)
  31. })
  32. }
  33. func (s *DatabaseSuite) TestCreateDatabase() {
  34. ctx, cancel := context.WithCancel(context.Background())
  35. defer cancel()
  36. s.Run("success", func() {
  37. dbName := fmt.Sprintf("dt_%s", s.randString(6))
  38. s.mock.EXPECT().CreateDatabase(mock.Anything, mock.Anything).RunAndReturn(func(ctx context.Context, cdr *milvuspb.CreateDatabaseRequest) (*commonpb.Status, error) {
  39. s.Equal(dbName, cdr.GetDbName())
  40. return merr.Success(), nil
  41. }).Once()
  42. err := s.client.CreateDatabase(ctx, NewCreateDatabaseOption(dbName))
  43. s.NoError(err)
  44. })
  45. s.Run("failure", func() {
  46. dbName := fmt.Sprintf("dt_%s", s.randString(6))
  47. s.mock.EXPECT().CreateDatabase(mock.Anything, mock.Anything).Return(nil, merr.WrapErrServiceInternal("mocked")).Once()
  48. err := s.client.CreateDatabase(ctx, NewCreateDatabaseOption(dbName))
  49. s.Error(err)
  50. })
  51. }
  52. func (s *DatabaseSuite) TestDropDatabase() {
  53. ctx, cancel := context.WithCancel(context.Background())
  54. defer cancel()
  55. s.Run("success", func() {
  56. dbName := fmt.Sprintf("dt_%s", s.randString(6))
  57. s.mock.EXPECT().DropDatabase(mock.Anything, mock.Anything).RunAndReturn(func(ctx context.Context, ddr *milvuspb.DropDatabaseRequest) (*commonpb.Status, error) {
  58. s.Equal(dbName, ddr.GetDbName())
  59. return merr.Success(), nil
  60. }).Once()
  61. err := s.client.DropDatabase(ctx, NewDropDatabaseOption(dbName))
  62. s.NoError(err)
  63. })
  64. s.Run("failure", func() {
  65. dbName := fmt.Sprintf("dt_%s", s.randString(6))
  66. s.mock.EXPECT().DropDatabase(mock.Anything, mock.Anything).Return(nil, merr.WrapErrServiceInternal("mocked")).Once()
  67. err := s.client.DropDatabase(ctx, NewDropDatabaseOption(dbName))
  68. s.Error(err)
  69. })
  70. }
  71. func TestDatabase(t *testing.T) {
  72. suite.Run(t, new(DatabaseSuite))
  73. }