milvus/client/database_test.go

package client

import (
	"context"
	"fmt"
	"testing"

	mock "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 DatabaseSuite struct {
	MockSuiteBase
}

func (s *DatabaseSuite) TestListDatabases() {
	ctx, cancel := context.WithCancel(context.Background())
	defer cancel()

	s.Run("success", func() {
		s.mock.EXPECT().ListDatabases(mock.Anything, mock.Anything).Return(&milvuspb.ListDatabasesResponse{
			Status:  merr.Success(),
			DbNames: []string{"default", "db1"},
		}, nil).Once()

		names, err := s.client.ListDatabase(ctx, NewListDatabaseOption())
		s.NoError(err)
		s.ElementsMatch([]string{"default", "db1"}, names)
	})

	s.Run("failure", func() {
		s.mock.EXPECT().ListDatabases(mock.Anything, mock.Anything).Return(nil, merr.WrapErrServiceInternal("mocked")).Once()

		_, err := s.client.ListDatabase(ctx, NewListDatabaseOption())
		s.Error(err)
	})
}

func (s *DatabaseSuite) TestCreateDatabase() {
	ctx, cancel := context.WithCancel(context.Background())
	defer cancel()

	s.Run("success", func() {
		dbName := fmt.Sprintf("dt_%s", s.randString(6))
		s.mock.EXPECT().CreateDatabase(mock.Anything, mock.Anything).RunAndReturn(func(ctx context.Context, cdr *milvuspb.CreateDatabaseRequest) (*commonpb.Status, error) {
			s.Equal(dbName, cdr.GetDbName())
			return merr.Success(), nil
		}).Once()

		err := s.client.CreateDatabase(ctx, NewCreateDatabaseOption(dbName))
		s.NoError(err)
	})

	s.Run("failure", func() {
		dbName := fmt.Sprintf("dt_%s", s.randString(6))
		s.mock.EXPECT().CreateDatabase(mock.Anything, mock.Anything).Return(nil, merr.WrapErrServiceInternal("mocked")).Once()

		err := s.client.CreateDatabase(ctx, NewCreateDatabaseOption(dbName))
		s.Error(err)
	})
}

func (s *DatabaseSuite) TestDropDatabase() {
	ctx, cancel := context.WithCancel(context.Background())
	defer cancel()

	s.Run("success", func() {
		dbName := fmt.Sprintf("dt_%s", s.randString(6))
		s.mock.EXPECT().DropDatabase(mock.Anything, mock.Anything).RunAndReturn(func(ctx context.Context, ddr *milvuspb.DropDatabaseRequest) (*commonpb.Status, error) {
			s.Equal(dbName, ddr.GetDbName())
			return merr.Success(), nil
		}).Once()

		err := s.client.DropDatabase(ctx, NewDropDatabaseOption(dbName))
		s.NoError(err)
	})

	s.Run("failure", func() {
		dbName := fmt.Sprintf("dt_%s", s.randString(6))
		s.mock.EXPECT().DropDatabase(mock.Anything, mock.Anything).Return(nil, merr.WrapErrServiceInternal("mocked")).Once()

		err := s.client.DropDatabase(ctx, NewDropDatabaseOption(dbName))
		s.Error(err)
	})
}

func TestDatabase(t *testing.T) {
	suite.Run(t, new(DatabaseSuite))
}