milvus/internal/metastore/db/dao/field_test.go

package dao

import (
	"errors"
	"testing"
	"time"

	"github.com/DATA-DOG/go-sqlmock"
	"github.com/milvus-io/milvus-proto/go-api/schemapb"
	"github.com/milvus-io/milvus/internal/metastore/db/dbmodel"
	"github.com/stretchr/testify/assert"
)

func TestField_GetByCollID(t *testing.T) {
	var fields = []*dbmodel.Field{
		{
			TenantID:     tenantID,
			FieldID:      fieldID1,
			FieldName:    "test_field_1",
			IsPrimaryKey: false,
			Description:  "",
			DataType:     schemapb.DataType_FloatVector,
			TypeParams:   "",
			IndexParams:  "",
			AutoID:       false,
			CollectionID: collID1,
			Ts:           ts,
		},
	}

	// expectation
	mock.ExpectQuery("SELECT * FROM `field_schemas` WHERE tenant_id = ? AND collection_id = ? AND ts = ? AND is_deleted = false").
		WithArgs(tenantID, collID1, ts).
		WillReturnRows(
			sqlmock.NewRows([]string{"tenant_id", "field_id", "field_name", "is_primary_key", "description", "data_type", "type_params", "index_params", "auto_id", "collection_id", "ts"}).
				AddRow(fields[0].TenantID, fields[0].FieldID, fields[0].FieldName, fields[0].IsPrimaryKey, fields[0].Description, fields[0].DataType, fields[0].TypeParams, fields[0].IndexParams, fields[0].AutoID, fields[0].CollectionID, fields[0].Ts))

	// actual
	res, err := fieldTestDb.GetByCollectionID(tenantID, collID1, ts)
	assert.Nil(t, err)
	assert.Equal(t, fields, res)
}

func TestField_GetByCollID_Error(t *testing.T) {
	// expectation
	mock.ExpectQuery("SELECT * FROM `field_schemas` WHERE tenant_id = ? AND collection_id = ? AND ts = ? AND is_deleted = false").
		WithArgs(tenantID, collID1, ts).
		WillReturnError(errors.New("test error"))

	// actual
	res, err := fieldTestDb.GetByCollectionID(tenantID, collID1, ts)
	assert.Nil(t, res)
	assert.Error(t, err)
}

func TestField_Insert(t *testing.T) {
	var fields = []*dbmodel.Field{
		{
			TenantID:     tenantID,
			FieldID:      fieldID1,
			FieldName:    "test_field_1",
			IsPrimaryKey: false,
			Description:  "",
			DataType:     schemapb.DataType_FloatVector,
			TypeParams:   "",
			IndexParams:  "",
			AutoID:       false,
			CollectionID: collID1,
			Ts:           ts,
			IsDeleted:    false,
			CreatedAt:    time.Now(),
			UpdatedAt:    time.Now(),
		},
	}

	// expectation
	mock.ExpectBegin()
	mock.ExpectExec("INSERT INTO `field_schemas` (`tenant_id`,`field_id`,`field_name`,`is_primary_key`,`description`,`data_type`,`type_params`,`index_params`,`auto_id`,`collection_id`,`ts`,`is_deleted`,`created_at`,`updated_at`) VALUES (?,?,?,?,?,?,?,?,?,?,?,?,?,?)").
		WithArgs(fields[0].TenantID, fields[0].FieldID, fields[0].FieldName, fields[0].IsPrimaryKey, fields[0].Description, fields[0].DataType, fields[0].TypeParams, fields[0].IndexParams, fields[0].AutoID, fields[0].CollectionID, fields[0].Ts, fields[0].IsDeleted, fields[0].CreatedAt, fields[0].UpdatedAt).
		WillReturnResult(sqlmock.NewResult(1, 1))
	mock.ExpectCommit()

	// actual
	err := fieldTestDb.Insert(fields)
	assert.Nil(t, err)
}

func TestField_Insert_Error(t *testing.T) {
	var fields = []*dbmodel.Field{
		{
			TenantID:     tenantID,
			FieldID:      fieldID1,
			FieldName:    "test_field_1",
			IsPrimaryKey: false,
			Description:  "",
			DataType:     schemapb.DataType_FloatVector,
			TypeParams:   "",
			IndexParams:  "",
			AutoID:       false,
			CollectionID: collID1,
			Ts:           ts,
			IsDeleted:    false,
			CreatedAt:    time.Now(),
			UpdatedAt:    time.Now(),
		},
	}

	// expectation
	mock.ExpectBegin()
	mock.ExpectExec("INSERT INTO `field_schemas` (`tenant_id`,`field_id`,`field_name`,`is_primary_key`,`description`,`data_type`,`type_params`,`index_params`,`auto_id`,`collection_id`,`ts`,`is_deleted`,`created_at`,`updated_at`) VALUES (?,?,?,?,?,?,?,?,?,?,?,?,?,?)").
		WithArgs(fields[0].TenantID, fields[0].FieldID, fields[0].FieldName, fields[0].IsPrimaryKey, fields[0].Description, fields[0].DataType, fields[0].TypeParams, fields[0].IndexParams, fields[0].AutoID, fields[0].CollectionID, fields[0].Ts, fields[0].IsDeleted, fields[0].CreatedAt, fields[0].UpdatedAt).
		WillReturnError(errors.New("test error"))
	mock.ExpectRollback()

	// actual
	err := fieldTestDb.Insert(fields)
	assert.Error(t, err)
}