[TST] Add tests for CheckCollections, make test collection setup more flexible (#4912)

## Summary
This adds a test for the `CheckCollections` method. This would have
caught the issue fixed in #4899.

Along the way, I discovered how these tests need some love, so I started
down this path by introducing a new `daotest` package that provides
testing helpers for creating `Collection` values that are easier to
configure. Instead of having to potentially debug all the tests that use
`dao.CreateTestCollection`, I defined a "shim" that can be used as a
stop-gap so that I can change the signature of this
`dao.CreateTestCollection` without potentially derailing the original
goal which was just to add some missing test coverage.

- Improvements & Bug fixes
  - Added test coverage for `CheckCollections` service endpoint
- New functionality
- Added `daotest` package and `NewTestCollection` helper with
builder/options methods for more configurability

## Test plan

All the tests are passing.

## Documentation Changes

N/A

Author: eculver

go/pkg/sysdb/grpc/collection_service_test.go

@@ -13,6 +13,7 @@ import (
 	"github.com/chroma-core/chroma/go/pkg/sysdb/coordinator"
 	"github.com/chroma-core/chroma/go/pkg/sysdb/coordinator/model"
 	"github.com/chroma-core/chroma/go/pkg/sysdb/metastore/db/dao"
+	"github.com/chroma-core/chroma/go/pkg/sysdb/metastore/db/dao/daotest"
 	"github.com/chroma-core/chroma/go/pkg/sysdb/metastore/db/dbcore"
 	s3metastore "github.com/chroma-core/chroma/go/pkg/sysdb/metastore/s3"
 	"github.com/chroma-core/chroma/go/pkg/types"
@@ -28,6 +29,11 @@ import (
 	"pgregory.net/rapid"
 )
 
+// TODO(eculver): replace most suite.NoError(err) with suite.Require().NoError(err) so the test
+// stops running when the error is not nil instead of continuing and causing red herrings in test output
+
+// TODO(eculver): replace calls to dao.NewDefaultTestCollection with daotest.NewTestCollection
+
 type CollectionServiceTestSuite struct {
 	suite.Suite
 	catalog      *coordinator.Catalog
@@ -346,9 +352,9 @@ func (suite *CollectionServiceTestSuite) TestCreateCollection() {
 }
 
 func (suite *CollectionServiceTestSuite) TestServer_GetCollection() {
-	// Create a test collection
+	// Create a test collection with a name that should not already exist in the database
 	collectionName := "test_get_collection"
-	collectionID, err := dao.CreateTestCollection(suite.db, collectionName, 128, suite.databaseId, nil)
+	collectionID, err := dao.CreateTestCollection(suite.db, daotest.NewDefaultTestCollection(collectionName, 128, suite.databaseId, nil))
 	suite.NoError(err)
 
 	// Soft delete the collection
@@ -377,6 +383,7 @@ func (suite *CollectionServiceTestSuite) TestServer_GetCollection() {
 
 func (suite *CollectionServiceTestSuite) TestServer_GetCollectionByResourceName() {
 	tenantResourceName := "test_tenant_resource_name"
+	// Does this need to match the daotest.TestTenantID?
 	tenantID := "test_tenant_id"
 	databaseName := "test_database"
 	collectionName := "test_collection"
@@ -388,7 +395,7 @@ func (suite *CollectionServiceTestSuite) TestServer_GetCollectionByResourceName(
 	err = dao.SetTestTenantResourceName(suite.db, tenantID, tenantResourceName)
 	suite.NoError(err)
 
-	collectionID, err := dao.CreateTestCollection(suite.db, collectionName, dim, databaseID, nil)
+	collectionID, err := dao.CreateTestCollection(suite.db, daotest.NewDefaultTestCollection(collectionName, dim, databaseID, nil))
 	suite.NoError(err)
 
 	req := &coordinatorpb.GetCollectionByResourceNameRequest{
@@ -447,7 +454,7 @@ func (suite *CollectionServiceTestSuite) TestServer_FlushCollectionCompaction()
 	log.Info("TestServer_FlushCollectionCompaction")
 	// create test collection
 	collectionName := "collection_service_test_flush_collection_compaction"
-	collectionID, err := dao.CreateTestCollection(suite.db, collectionName, 128, suite.databaseId, nil)
+	collectionID, err := dao.CreateTestCollection(suite.db, daotest.NewDefaultTestCollection(collectionName, 128, suite.databaseId, nil))
 	suite.NoError(err)
 
 	// flush collection compaction
@@ -597,7 +604,7 @@ func (suite *CollectionServiceTestSuite) TestServer_FlushCollectionCompaction()
 	// Send FlushCollectionCompaction for a collection that is soft deleted.
 	// It should fail with a failed precondition error.
 	// Create collection and soft-delete it.
-	collectionID, err = dao.CreateTestCollection(suite.db, "test_flush_collection_compaction_soft_delete", 128, suite.databaseId, nil)
+	collectionID, err = dao.CreateTestCollection(suite.db, daotest.NewDefaultTestCollection("test_flush_collection_compaction_soft_delete", 128, suite.databaseId, nil))
 	suite.NoError(err)
 	suite.s.coordinator.SoftDeleteCollection(context.Background(), &model.DeleteCollection{
 		ID:           types.MustParse(collectionID),
@@ -621,9 +628,35 @@ func (suite *CollectionServiceTestSuite) TestServer_FlushCollectionCompaction()
 	suite.NoError(err)
 }
 
+func (suite *CollectionServiceTestSuite) TestServer_CheckCollections() {
+	collectionName := "test_check_collections"
+	collectionID, err := dao.CreateTestCollection(suite.db, daotest.NewDefaultTestCollection(collectionName, 128, suite.databaseId, nil))
+	suite.NoError(err)
+
+	request := &coordinatorpb.CheckCollectionsRequest{
+		CollectionIds: []string{collectionID},
+	}
+
+	// Call the service method
+	response, err := suite.s.CheckCollections(context.Background(), request)
+	suite.NoError(err)
+
+	suite.NotNil(response.GetDeleted(), "Deleted slice should not be nil.")
+	suite.Len(response.GetDeleted(), 1)
+	suite.False(response.GetDeleted()[0])
+
+	suite.NotNil(response.GetLogPosition(), "LogPosition slice should not be nil.")
+	suite.Len(response.GetLogPosition(), 1)
+	suite.GreaterOrEqual(response.GetLogPosition()[0], int64(0))
+
+	// clean up
+	err = dao.CleanUpTestCollection(suite.db, collectionID)
+	suite.NoError(err)
+}
+
 func (suite *CollectionServiceTestSuite) TestGetCollectionSize() {
 	collectionName := "collection_service_test_get_collection_size"
-	collectionID, err := dao.CreateTestCollection(suite.db, collectionName, 128, suite.databaseId, nil)
+	collectionID, err := dao.CreateTestCollection(suite.db, daotest.NewDefaultTestCollection(collectionName, 128, suite.databaseId, nil))
 	suite.NoError(err)
 
 	req := coordinatorpb.GetCollectionSizeRequest{
@@ -639,7 +672,7 @@ func (suite *CollectionServiceTestSuite) TestGetCollectionSize() {
 
 func (suite *CollectionServiceTestSuite) TestCountForks() {
 	collectionName := "collection_service_test_count_forks"
-	collectionID, err := dao.CreateTestCollection(suite.db, collectionName, 128, suite.databaseId, nil)
+	collectionID, err := dao.CreateTestCollection(suite.db, daotest.NewDefaultTestCollection(collectionName, 128, suite.databaseId, nil))
 	suite.NoError(err)
 
 	req := coordinatorpb.CountForksRequest{
@@ -684,7 +717,7 @@ func (suite *CollectionServiceTestSuite) TestCountForks() {
 
 func (suite *CollectionServiceTestSuite) TestFork() {
 	collectionName := "collection_service_test_forks"
-	collectionID, err := dao.CreateTestCollection(suite.db, collectionName, 128, suite.databaseId, nil)
+	collectionID, err := dao.CreateTestCollection(suite.db, daotest.NewDefaultTestCollection(collectionName, 128, suite.databaseId, nil))
 	suite.NoError(err)
 	targetCollectionID := types.NewUniqueID()
 

go/pkg/sysdb/metastore/db/dao/collection_test.go

@@ -5,6 +5,7 @@ import (
 	"testing"
 	"time"
 
+	"github.com/chroma-core/chroma/go/pkg/sysdb/metastore/db/dao/daotest"
 	"github.com/chroma-core/chroma/go/pkg/sysdb/metastore/db/dbcore"
 	"github.com/pingcap/log"
 	"github.com/stretchr/testify/suite"
@@ -48,7 +49,7 @@ func (suite *CollectionDbTestSuite) TearDownSuite() {
 func (suite *CollectionDbTestSuite) TestCollectionDb_GetCollections() {
 	collectionName := "test_collection_get_collections"
 	dim := int32(128)
-	collectionID, err := CreateTestCollection(suite.db, collectionName, dim, suite.databaseId, nil)
+	collectionID, err := CreateTestCollection(suite.db, daotest.NewDefaultTestCollection(collectionName, dim, suite.databaseId, nil))
 	suite.NoError(err)
 
 	testKey := "test"
@@ -103,7 +104,7 @@ func (suite *CollectionDbTestSuite) TestCollectionDb_GetCollections() {
 	suite.Len(collections, 1)
 	suite.Equal(collectionID, collections[0].Collection.ID)
 
-	collectionID2, err := CreateTestCollection(suite.db, "test_collection_get_collections2", 128, suite.databaseId, nil)
+	collectionID2, err := CreateTestCollection(suite.db, daotest.NewDefaultTestCollection("test_collection_get_collections2", 128, suite.databaseId, nil))
 	suite.NoError(err)
 
 	// Test order by. Collections are ordered by create time so collectionID2 should be second
@@ -137,10 +138,10 @@ func (suite *CollectionDbTestSuite) TestCollectionDb_GetCollections() {
 	suite.NoError(err)
 
 	// Create two collections in the new database.
-	collectionID3, err := CreateTestCollection(suite.db, "test_collection_get_collections3", 128, DbId, nil)
+	collectionID3, err := CreateTestCollection(suite.db, daotest.NewDefaultTestCollection("test_collection_get_collections3", 128, DbId, nil))
 	suite.NoError(err)
 
-	collectionID4, err := CreateTestCollection(suite.db, "test_collection_get_collections4", 128, DbId, nil)
+	collectionID4, err := CreateTestCollection(suite.db, daotest.NewDefaultTestCollection("test_collection_get_collections4", 128, DbId, nil))
 	suite.NoError(err)
 
 	// Test count collections
@@ -174,7 +175,7 @@ func (suite *CollectionDbTestSuite) TestCollectionDb_GetCollections() {
 
 func (suite *CollectionDbTestSuite) TestCollectionDb_UpdateLogPositionVersionTotalRecordsAndLogicalSize() {
 	collectionName := "test_collection_get_collections"
-	collectionID, _ := CreateTestCollection(suite.db, collectionName, 128, suite.databaseId, nil)
+	collectionID, _ := CreateTestCollection(suite.db, daotest.NewDefaultTestCollection(collectionName, 128, suite.databaseId, nil))
 	ids := []string{collectionID}
 	// verify default values
 	collections, err := suite.collectionDb.GetCollections(ids, nil, "", "", nil, nil, false)
@@ -228,9 +229,9 @@ func (suite *CollectionDbTestSuite) TestCollectionDb_SoftDelete() {
 	// Create 2 collections.
 	collectionName1 := "test_collection_soft_delete1"
 	collectionName2 := "test_collection_soft_delete2"
-	collectionID1, err := CreateTestCollection(suite.db, collectionName1, 128, suite.databaseId, nil)
+	collectionID1, err := CreateTestCollection(suite.db, daotest.NewDefaultTestCollection(collectionName1, 128, suite.databaseId, nil))
 	suite.NoError(err)
-	collectionID2, err := CreateTestCollection(suite.db, collectionName2, 128, suite.databaseId, nil)
+	collectionID2, err := CreateTestCollection(suite.db, daotest.NewDefaultTestCollection(collectionName2, 128, suite.databaseId, nil))
 	suite.NoError(err)
 
 	// Soft delete collection 1 by Updating the is_deleted column
@@ -265,7 +266,7 @@ func (suite *CollectionDbTestSuite) TestCollectionDb_SoftDelete() {
 
 func (suite *CollectionDbTestSuite) TestCollectionDb_GetCollectionSize() {
 	collectionName := "test_collection_get_collection_size"
-	collectionID, err := CreateTestCollection(suite.db, collectionName, 128, suite.databaseId, nil)
+	collectionID, err := CreateTestCollection(suite.db, daotest.NewDefaultTestCollection(collectionName, 128, suite.databaseId, nil))
 	suite.NoError(err)
 
 	total_records_post_compaction, err := suite.collectionDb.GetCollectionSize(collectionID)
@@ -299,17 +300,17 @@ func (suite *CollectionDbTestSuite) TestCollectionDb_GetCollectionByResourceName
 
 	collectionName := "test_collection"
 	dim := int32(128)
-	collectionID, err := CreateTestCollection(suite.db, collectionName, dim, databaseID, nil)
+	collectionID, err := CreateTestCollection(suite.db, daotest.NewDefaultTestCollection(collectionName, dim, databaseID, nil))
 	suite.NoError(err)
 
-	collection, err := suite.collectionDb.GetCollectionByResourceName(tenantResourceName, databaseName, collectionName)
+	collectionResult, err := suite.collectionDb.GetCollectionByResourceName(tenantResourceName, databaseName, collectionName)
 	suite.NoError(err)
-	suite.NotNil(collection)
-	suite.Equal(collectionID, collection.Collection.ID)
-	suite.Equal(collectionName, *collection.Collection.Name)
-	suite.Equal(databaseID, collection.Collection.DatabaseID)
-	suite.Equal(tenantID, collection.TenantID)
-	suite.Equal(databaseName, collection.DatabaseName)
+	suite.NotNil(collectionResult)
+	suite.Equal(collectionID, collectionResult.Collection.ID)
+	suite.Equal(collectionName, *collectionResult.Collection.Name)
+	suite.Equal(databaseID, collectionResult.Collection.DatabaseID)
+	suite.Equal(tenantID, collectionResult.TenantID)
+	suite.Equal(databaseName, collectionResult.DatabaseName)
 
 	nonExistentCollection, err := suite.collectionDb.GetCollectionByResourceName(tenantResourceName, databaseName, "non_existent_collection")
 	suite.Error(err, "collection not found")

go/pkg/sysdb/metastore/db/dao/daotest/collection.go

@@ -0,0 +1,166 @@
+package daotest
+
+import (
+	"time"
+
+	"github.com/chroma-core/chroma/go/pkg/sysdb/metastore/db/dbmodel"
+	"github.com/chroma-core/chroma/go/pkg/types"
+	"github.com/pingcap/log"
+	"go.uber.org/zap"
+)
+
+// defaults for a test collection
+var (
+	defaultConfigurationJsonStr       = "{\"a\": \"param\", \"b\": \"param2\", \"3\": true}"
+	defaultDimension                  = int32(128)
+	defaultTotalRecordsPostCompaction = uint64(100)
+	defaultSizeBytesPostCompaction    = uint64(500000)
+	defaultLastCompactionTimeSecs     = uint64(1741037006)
+)
+
+// TestDatabaseID is the default database ID for a test collection. It's exported because it's so frequently used
+// and can be the same across most tests.
+const TestDatabaseID = "test_database"
+
+// TestTenantID is the default tenant ID for a test collection. It's exported because it's so frequently used
+// and can be the same across most tests.
+const TestTenantID = "test_tenant"
+
+// NewDefaultTestCollection is a "shim" for callsites that existed before dao.CreateTestCollection was refactored to
+// take a dbmodel.Collection instead of a subset of its fields as arguments. It should not be used for new tests.
+func NewDefaultTestCollection(collectionName string, dimension int32, databaseID string, lineageFileName *string) *dbmodel.Collection {
+	log.Info("new default test collection", zap.String("collectionName", collectionName), zap.Int32("dimension", dimension), zap.String("databaseID", databaseID))
+	return NewTestCollection(
+		TestTenantID,
+		databaseID,
+		collectionName,
+		WithDimension(dimension),
+		WithTotalRecordsPostCompaction(defaultTotalRecordsPostCompaction),
+		WithSizeBytesPostCompaction(defaultSizeBytesPostCompaction),
+		WithLastCompactionTimeSecs(defaultLastCompactionTimeSecs),
+		WithLineageFileName(lineageFileName),
+	)
+}
+
+// NewTestCollection creates a new test collection with the given name, database ID, and tenant ID.
+// Name, databaseID, and tenantID are required, other fields have defaults but can be overridden with
+// option function of a similar name.
+// Note: collection.CreatedAt is set to the current time, but collection.UpdatedAt is not set.
+func NewTestCollection(tenantID, databaseID, collectionName string, options ...TestCollectionOption) *dbmodel.Collection {
+	log.Info("new test collection", zap.String("tenantID", tenantID), zap.String("databaseID", databaseID), zap.String("collectionName", collectionName))
+	collectionId := types.NewUniqueID().String()
+
+	collection := &dbmodel.Collection{
+		ID:                   collectionId,
+		Name:                 &collectionName,
+		ConfigurationJsonStr: &defaultConfigurationJsonStr,
+		Dimension:            &defaultDimension,
+		DatabaseID:           databaseID,
+		CreatedAt:            time.Now(),
+		Tenant:               tenantID,
+	}
+
+	for _, option := range options {
+		option(collection)
+	}
+
+	return collection
+}
+
+type TestCollectionOption func(*dbmodel.Collection)
+
+func WithConfigurationJsonStr(configurationJsonStr string) func(*dbmodel.Collection) {
+	return func(collection *dbmodel.Collection) {
+		collection.ConfigurationJsonStr = &configurationJsonStr
+	}
+}
+
+func WithDimension(dimension int32) func(*dbmodel.Collection) {
+	return func(collection *dbmodel.Collection) {
+		collection.Dimension = &dimension
+	}
+}
+
+func WithTs(ts int64) func(*dbmodel.Collection) {
+	return func(collection *dbmodel.Collection) {
+		collection.Ts = ts
+	}
+}
+
+func WithIsDeleted(isDeleted bool) func(*dbmodel.Collection) {
+	return func(collection *dbmodel.Collection) {
+		collection.IsDeleted = isDeleted
+	}
+}
+
+func WithCreatedAt(createdAt time.Time) func(*dbmodel.Collection) {
+	return func(collection *dbmodel.Collection) {
+		collection.CreatedAt = createdAt
+	}
+}
+
+func WithUpdatedAt(updatedAt time.Time) func(*dbmodel.Collection) {
+	return func(collection *dbmodel.Collection) {
+		collection.UpdatedAt = updatedAt
+	}
+}
+
+func WithLogPosition(logPosition int64) func(*dbmodel.Collection) {
+	return func(collection *dbmodel.Collection) {
+		collection.LogPosition = logPosition
+	}
+}
+
+func WithVersion(version int32) func(*dbmodel.Collection) {
+	return func(collection *dbmodel.Collection) {
+		collection.Version = version
+	}
+}
+
+func WithVersionFileName(versionFileName string) func(*dbmodel.Collection) {
+	return func(collection *dbmodel.Collection) {
+		collection.VersionFileName = versionFileName
+	}
+}
+
+func WithRootCollectionID(rootCollectionId string) func(*dbmodel.Collection) {
+	return func(collection *dbmodel.Collection) {
+		collection.RootCollectionId = &rootCollectionId
+	}
+}
+
+func WithLineageFileName(lineageFileName *string) func(*dbmodel.Collection) {
+	return func(collection *dbmodel.Collection) {
+		collection.LineageFileName = lineageFileName
+	}
+}
+
+func WithTotalRecordsPostCompaction(totalRecordsPostCompaction uint64) func(*dbmodel.Collection) {
+	return func(collection *dbmodel.Collection) {
+		collection.TotalRecordsPostCompaction = totalRecordsPostCompaction
+	}
+}
+
+func WithSizeBytesPostCompaction(sizeBytesPostCompaction uint64) func(*dbmodel.Collection) {
+	return func(collection *dbmodel.Collection) {
+		collection.SizeBytesPostCompaction = sizeBytesPostCompaction
+	}
+}
+
+func WithLastCompactionTimeSecs(lastCompactionTimeSecs uint64) func(*dbmodel.Collection) {
+	return func(collection *dbmodel.Collection) {
+		collection.LastCompactionTimeSecs = lastCompactionTimeSecs
+	}
+}
+
+func WithNumVersions(numVersions uint32) func(*dbmodel.Collection) {
+	return func(collection *dbmodel.Collection) {
+		collection.NumVersions = numVersions
+	}
+}
+
+func WithOldestVersionTs(oldestVersionTs time.Time) func(*dbmodel.Collection) {
+	return func(collection *dbmodel.Collection) {
+		collection.OldestVersionTs = oldestVersionTs
+	}
+}