Skip to content

Commit bbf5511

Browse files
sanketkediasanketkedia
committed
Fix tests
1 parent d5d9841 commit bbf5511

File tree

7 files changed

+141
-35
lines changed

7 files changed

+141
-35
lines changed

‎go/pkg/sysdb/coordinator/coordinator.go‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -131,7 +131,7 @@ func (s *Coordinator) GetCollectionSize(ctx context.Context, collectionID types.
131131
}
132132

133133
func (s *Coordinator) GetCollectionWithSegments(ctx context.Context, collectionID types.UniqueID) (*model.Collection, []*model.Segment, error) {
134-
return s.catalog.GetCollectionWithSegments(ctx, collectionID)
134+
return s.catalog.GetCollectionWithSegments(ctx, collectionID, false)
135135
}
136136

137137
func (s *Coordinator) CheckCollection(ctx context.Context, collectionID types.UniqueID) (bool, error) {

‎go/pkg/sysdb/coordinator/coordinator_test.go‎

Lines changed: 109 additions & 16 deletions
Original file line numberDiff line numberDiff line change
@@ -407,7 +407,15 @@ func (suite *APIsTestSuite) TestCreateGetDeleteCollections() {
407407
sort.Slice(results, func(i, j int) bool {
408408
return results[i].Name < results[j].Name
409409
})
410-
suite.Equal(suite.sampleCollections, results)
410+
suite.Len(results, len(suite.sampleCollections))
411+
for i, collection := range results {
412+
suite.Equal(suite.sampleCollections[i].ID, collection.ID)
413+
suite.Equal(suite.sampleCollections[i].Name, collection.Name)
414+
suite.Equal(suite.sampleCollections[i].TenantID, collection.TenantID)
415+
suite.Equal(suite.sampleCollections[i].DatabaseName, collection.DatabaseName)
416+
suite.Equal(suite.sampleCollections[i].Dimension, collection.Dimension)
417+
suite.Equal(suite.sampleCollections[i].Metadata, collection.Metadata)
418+
}
411419

412420
// Duplicate create fails
413421
_, _, err = suite.coordinator.CreateCollection(ctx, &model.CreateCollection{
@@ -448,14 +456,26 @@ func (suite *APIsTestSuite) TestCreateGetDeleteCollections() {
448456
for _, collection := range suite.sampleCollections {
449457
result, err := suite.coordinator.GetCollections(ctx, types.NilUniqueID(), &collection.Name, suite.tenantName, suite.databaseName, nil, nil)
450458
suite.NoError(err)
451-
suite.Equal([]*model.Collection{collection}, result)
459+
suite.Equal(len(result), 1)
460+
suite.Equal(collection.ID, result[0].ID)
461+
suite.Equal(collection.Name, result[0].Name)
462+
suite.Equal(collection.TenantID, result[0].TenantID)
463+
suite.Equal(collection.DatabaseName, result[0].DatabaseName)
464+
suite.Equal(collection.Dimension, result[0].Dimension)
465+
suite.Equal(collection.Metadata, result[0].Metadata)
452466
}
453467

454468
// Find by id
455469
for _, collection := range suite.sampleCollections {
456470
result, err := suite.coordinator.GetCollections(ctx, collection.ID, nil, suite.tenantName, suite.databaseName, nil, nil)
457471
suite.NoError(err)
458-
suite.Equal([]*model.Collection{collection}, result)
472+
suite.Equal(len(result), 1)
473+
suite.Equal(collection.ID, result[0].ID)
474+
suite.Equal(collection.Name, result[0].Name)
475+
suite.Equal(collection.TenantID, result[0].TenantID)
476+
suite.Equal(collection.DatabaseName, result[0].DatabaseName)
477+
suite.Equal(collection.Dimension, result[0].Dimension)
478+
suite.Equal(collection.Metadata, result[0].Metadata)
459479
}
460480

461481
// Delete
@@ -470,10 +490,18 @@ func (suite *APIsTestSuite) TestCreateGetDeleteCollections() {
470490

471491
results, err = suite.coordinator.GetCollections(ctx, types.NilUniqueID(), nil, suite.tenantName, suite.databaseName, nil, nil)
472492
suite.NoError(err)
493+
result_ids := make([]types.UniqueID, len(results))
494+
for i, result := range results {
495+
result_ids[i] = result.ID
496+
}
473497

474-
suite.NotContains(results, c1)
498+
sample_ids := make([]types.UniqueID, len(suite.sampleCollections))
499+
for i, collection := range suite.sampleCollections {
500+
sample_ids[i] = collection.ID
501+
}
502+
suite.NotContains(result_ids, c1.ID)
475503
suite.Len(results, len(suite.sampleCollections)-1)
476-
suite.ElementsMatch(results, suite.sampleCollections[1:])
504+
suite.ElementsMatch(result_ids, sample_ids[1:])
477505
byIDResult, err := suite.coordinator.GetCollections(ctx, c1.ID, nil, suite.tenantName, suite.databaseName, nil, nil)
478506
suite.NoError(err)
479507
suite.Empty(byIDResult)
@@ -632,40 +660,88 @@ func (suite *APIsTestSuite) TestUpdateCollections() {
632660
coll.Name = "new_name"
633661
result, err := suite.coordinator.UpdateCollection(ctx, &model.UpdateCollection{ID: coll.ID, Name: &coll.Name})
634662
suite.NoError(err)
635-
suite.Equal(coll, result)
663+
suite.Equal(coll.ID, result.ID)
664+
suite.Equal(coll.Name, result.Name)
665+
suite.Equal(coll.TenantID, result.TenantID)
666+
suite.Equal(coll.DatabaseName, result.DatabaseName)
667+
suite.Equal(coll.Dimension, result.Dimension)
668+
suite.Equal(coll.Metadata, result.Metadata)
669+
636670
resultList, err := suite.coordinator.GetCollections(ctx, types.NilUniqueID(), &coll.Name, suite.tenantName, suite.databaseName, nil, nil)
637671
suite.NoError(err)
638-
suite.Equal([]*model.Collection{coll}, resultList)
672+
suite.Equal(len(resultList), 1)
673+
suite.Equal(coll.ID, resultList[0].ID)
674+
suite.Equal(coll.Name, resultList[0].Name)
675+
suite.Equal(coll.TenantID, resultList[0].TenantID)
676+
suite.Equal(coll.DatabaseName, resultList[0].DatabaseName)
677+
suite.Equal(coll.Dimension, resultList[0].Dimension)
678+
suite.Equal(coll.Metadata, resultList[0].Metadata)
639679

640680
// Update dimension
641681
newDimension := int32(128)
642682
coll.Dimension = &newDimension
643683
result, err = suite.coordinator.UpdateCollection(ctx, &model.UpdateCollection{ID: coll.ID, Dimension: coll.Dimension})
644684
suite.NoError(err)
645-
suite.Equal(coll, result)
685+
suite.Equal(coll.ID, result.ID)
686+
suite.Equal(coll.Name, result.Name)
687+
suite.Equal(coll.TenantID, result.TenantID)
688+
suite.Equal(coll.DatabaseName, result.DatabaseName)
689+
suite.Equal(coll.Dimension, result.Dimension)
690+
suite.Equal(coll.Metadata, result.Metadata)
691+
646692
resultList, err = suite.coordinator.GetCollections(ctx, coll.ID, nil, suite.tenantName, suite.databaseName, nil, nil)
647693
suite.NoError(err)
648-
suite.Equal([]*model.Collection{coll}, resultList)
694+
suite.Equal(len(resultList), 1)
695+
suite.Equal(coll.ID, resultList[0].ID)
696+
suite.Equal(coll.Name, resultList[0].Name)
697+
suite.Equal(coll.TenantID, resultList[0].TenantID)
698+
suite.Equal(coll.DatabaseName, resultList[0].DatabaseName)
699+
suite.Equal(coll.Dimension, resultList[0].Dimension)
700+
suite.Equal(coll.Metadata, resultList[0].Metadata)
649701

650702
// Reset the metadata
651703
newMetadata := model.NewCollectionMetadata[model.CollectionMetadataValueType]()
652704
newMetadata.Add("test_str2", &model.CollectionMetadataValueStringType{Value: "str2"})
653705
coll.Metadata = newMetadata
654706
result, err = suite.coordinator.UpdateCollection(ctx, &model.UpdateCollection{ID: coll.ID, Metadata: coll.Metadata})
655707
suite.NoError(err)
656-
suite.Equal(coll, result)
708+
suite.Equal(coll.ID, result.ID)
709+
suite.Equal(coll.Name, result.Name)
710+
suite.Equal(coll.TenantID, result.TenantID)
711+
suite.Equal(coll.DatabaseName, result.DatabaseName)
712+
suite.Equal(coll.Dimension, result.Dimension)
713+
suite.Equal(coll.Metadata, result.Metadata)
714+
657715
resultList, err = suite.coordinator.GetCollections(ctx, coll.ID, nil, suite.tenantName, suite.databaseName, nil, nil)
658716
suite.NoError(err)
659-
suite.Equal([]*model.Collection{coll}, resultList)
717+
suite.Equal(len(resultList), 1)
718+
suite.Equal(coll.ID, resultList[0].ID)
719+
suite.Equal(coll.Name, resultList[0].Name)
720+
suite.Equal(coll.TenantID, resultList[0].TenantID)
721+
suite.Equal(coll.DatabaseName, resultList[0].DatabaseName)
722+
suite.Equal(coll.Dimension, resultList[0].Dimension)
723+
suite.Equal(coll.Metadata, resultList[0].Metadata)
660724

661725
// Delete all metadata keys
662726
coll.Metadata = nil
663727
result, err = suite.coordinator.UpdateCollection(ctx, &model.UpdateCollection{ID: coll.ID, Metadata: coll.Metadata, ResetMetadata: true})
664728
suite.NoError(err)
665-
suite.Equal(coll, result)
729+
suite.Equal(coll.ID, result.ID)
730+
suite.Equal(coll.Name, result.Name)
731+
suite.Equal(coll.TenantID, result.TenantID)
732+
suite.Equal(coll.DatabaseName, result.DatabaseName)
733+
suite.Equal(coll.Dimension, result.Dimension)
734+
suite.Equal(coll.Metadata, result.Metadata)
735+
666736
resultList, err = suite.coordinator.GetCollections(ctx, coll.ID, nil, suite.tenantName, suite.databaseName, nil, nil)
667737
suite.NoError(err)
668-
suite.Equal([]*model.Collection{coll}, resultList)
738+
suite.Equal(len(resultList), 1)
739+
suite.Equal(coll.ID, resultList[0].ID)
740+
suite.Equal(coll.Name, resultList[0].Name)
741+
suite.Equal(coll.TenantID, resultList[0].TenantID)
742+
suite.Equal(coll.DatabaseName, resultList[0].DatabaseName)
743+
suite.Equal(coll.Dimension, resultList[0].Dimension)
744+
suite.Equal(coll.Metadata, resultList[0].Metadata)
669745
}
670746

671747
func (suite *APIsTestSuite) TestGetOrCreateCollectionsTwice() {
@@ -794,7 +870,14 @@ func (suite *APIsTestSuite) TestGetMultipleWithDatabase() {
794870
sort.Slice(result, func(i, j int) bool {
795871
return result[i].Name < result[j].Name
796872
})
797-
suite.Equal(suite.sampleCollections, result)
873+
for index, collection := range result {
874+
suite.Equal(suite.sampleCollections[index].ID, collection.ID)
875+
suite.Equal(suite.sampleCollections[index].Name, collection.Name)
876+
suite.Equal(suite.sampleCollections[index].TenantID, collection.TenantID)
877+
suite.Equal(suite.sampleCollections[index].DatabaseName, collection.DatabaseName)
878+
suite.Equal(suite.sampleCollections[index].Dimension, collection.Dimension)
879+
suite.Equal(suite.sampleCollections[index].Metadata, collection.Metadata)
880+
}
798881

799882
result, err = suite.coordinator.GetCollections(ctx, types.NilUniqueID(), nil, suite.tenantName, suite.databaseName, nil, nil)
800883
suite.NoError(err)
@@ -876,15 +959,25 @@ func (suite *APIsTestSuite) TestCreateDatabaseWithTenants() {
876959
result, err := suite.coordinator.GetCollections(ctx, types.NilUniqueID(), nil, newTenantName, newDatabaseName, nil, nil)
877960
suite.NoError(err)
878961
suite.Len(result, 1)
879-
suite.Equal(expected[0], result[0])
962+
suite.Equal(expected[0].ID, result[0].ID)
963+
suite.Equal(expected[0].Name, result[0].Name)
964+
suite.Equal(expected[0].TenantID, result[0].TenantID)
965+
suite.Equal(expected[0].DatabaseName, result[0].DatabaseName)
966+
suite.Equal(expected[0].Dimension, result[0].Dimension)
967+
suite.Equal(expected[0].Metadata, result[0].Metadata)
880968

881969
expected = []*model.Collection{suite.sampleCollections[1]}
882970
expected[0].TenantID = suite.tenantName
883971
expected[0].DatabaseName = newDatabaseName
884972
result, err = suite.coordinator.GetCollections(ctx, types.NilUniqueID(), nil, suite.tenantName, newDatabaseName, nil, nil)
885973
suite.NoError(err)
886974
suite.Len(result, 1)
887-
suite.Equal(expected[0], result[0])
975+
suite.Equal(expected[0].ID, result[0].ID)
976+
suite.Equal(expected[0].Name, result[0].Name)
977+
suite.Equal(expected[0].TenantID, result[0].TenantID)
978+
suite.Equal(expected[0].DatabaseName, result[0].DatabaseName)
979+
suite.Equal(expected[0].Dimension, result[0].Dimension)
980+
suite.Equal(expected[0].Metadata, result[0].Metadata)
888981

889982
// A new tenant DOES NOT have a default database. This does not error, instead 0
890983
// results are returned

‎go/pkg/sysdb/coordinator/model_db_convert_test.go‎

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -138,6 +138,7 @@ func TestConvertCollectionToModel(t *testing.T) {
138138
collectionTotalRecordsPostCompaction := uint64(100)
139139
collectionSizeBytesPostCompaction := uint64(500000)
140140
collectionLastCompactionTimeSecs := uint64(1741037006)
141+
dbId := types.NewUniqueID()
141142
collectionAndMetadata := &dbmodel.CollectionAndMetadata{
142143
Collection: &dbmodel.Collection{
143144
ID: collectionID.String(),
@@ -147,6 +148,7 @@ func TestConvertCollectionToModel(t *testing.T) {
147148
TotalRecordsPostCompaction: collectionTotalRecordsPostCompaction,
148149
SizeBytesPostCompaction: collectionSizeBytesPostCompaction,
149150
LastCompactionTimeSecs: collectionLastCompactionTimeSecs,
151+
DatabaseID: dbId.String(),
150152
},
151153
CollectionMetadata: []*dbmodel.CollectionMetadata{},
152154
}
@@ -161,4 +163,5 @@ func TestConvertCollectionToModel(t *testing.T) {
161163
assert.Equal(t, collectionSizeBytesPostCompaction, modelCollections[0].SizeBytesPostCompaction)
162164
assert.Equal(t, collectionLastCompactionTimeSecs, modelCollections[0].LastCompactionTimeSecs)
163165
assert.Nil(t, modelCollections[0].Metadata)
166+
assert.Equal(t, dbId, modelCollections[0].DatabaseId)
164167
}

‎go/pkg/sysdb/coordinator/table_catalog.go‎

Lines changed: 9 additions & 11 deletions
Original file line numberDiff line numberDiff line change
@@ -460,7 +460,7 @@ func (tc *Catalog) ListCollectionsToGc(ctx context.Context, cutoffTimeSecs *uint
460460
return collections, nil
461461
}
462462

463-
func (tc *Catalog) GetCollectionWithSegments(ctx context.Context, collectionID types.UniqueID) (*model.Collection, []*model.Segment, error) {
463+
func (tc *Catalog) GetCollectionWithSegments(ctx context.Context, collectionID types.UniqueID, returnSoftDeleted bool) (*model.Collection, []*model.Segment, error) {
464464
tracer := otel.Tracer
465465
if tracer != nil {
466466
_, span := tracer.Start(ctx, "Catalog.GetCollections")
@@ -471,23 +471,23 @@ func (tc *Catalog) GetCollectionWithSegments(ctx context.Context, collectionID t
471471
var segments []*model.Segment
472472

473473
err := tc.txImpl.Transaction(ctx, func(txCtx context.Context) error {
474-
collections, e := tc.GetCollections(txCtx, collectionID, nil, "", "", nil, nil)
474+
collection_entry, e := tc.GetCollection(txCtx, collectionID, nil, "", "")
475475
if e != nil {
476476
return e
477477
}
478-
if len(collections) == 0 {
478+
if collection_entry == nil {
479479
return common.ErrCollectionNotFound
480480
}
481-
if len(collections) > 1 {
482-
return common.ErrCollectionUniqueConstraintViolation
481+
if collection_entry.IsDeleted && !returnSoftDeleted {
482+
return common.ErrCollectionNotFound
483483
}
484-
collection = collections[0]
485484

486485
segments, e = tc.GetSegments(txCtx, types.NilUniqueID(), nil, nil, collectionID)
487486
if e != nil {
488487
return e
489488
}
490489

490+
collection = collection_entry
491491
return nil
492492
})
493493
if err != nil {
@@ -876,7 +876,7 @@ func (tc *Catalog) ForkCollection(ctx context.Context, forkCollection *model.For
876876
}
877877

878878
// Get source and root collections after they are locked
879-
sourceCollection, sourceSegments, err = tc.GetCollectionWithSegments(txCtx, forkCollection.SourceCollectionID)
879+
sourceCollection, sourceSegments, err = tc.GetCollectionWithSegments(txCtx, forkCollection.SourceCollectionID, false)
880880
if err != nil {
881881
return err
882882
}
@@ -990,7 +990,7 @@ func (tc *Catalog) ForkCollection(ctx context.Context, forkCollection *model.For
990990
return nil, nil, err
991991
}
992992

993-
return tc.GetCollectionWithSegments(ctx, forkCollection.TargetCollectionID)
993+
return tc.GetCollectionWithSegments(ctx, forkCollection.TargetCollectionID, false)
994994
}
995995

996996
func (tc *Catalog) CountForks(ctx context.Context, sourceCollectionID types.UniqueID) (uint64, error) {
@@ -1577,8 +1577,7 @@ func (tc *Catalog) FlushCollectionCompactionForVersionedCollection(ctx context.C
15771577
for numAttempts < maxAttempts {
15781578
numAttempts++
15791579
// Get the current version info and the version file from the table.
1580-
collectionEntry, segments, err := tc.GetCollectionWithSegments(ctx, flushCollectionCompaction.ID)
1581-
// collectionEntry, err := tc.metaDomain.CollectionDb(ctx).GetCollectionEntry(types.FromUniqueID(flushCollectionCompaction.ID), nil)
1580+
collectionEntry, segments, err := tc.GetCollectionWithSegments(ctx, flushCollectionCompaction.ID, true)
15821581
if err != nil {
15831582
return nil, err
15841583
}
@@ -1608,7 +1607,6 @@ func (tc *Catalog) FlushCollectionCompactionForVersionedCollection(ctx context.C
16081607
}
16091608

16101609
existingVersionFileName := collectionEntry.VersionFileName
1611-
// existingSegments := segments
16121610
var existingVersionFilePb *coordinatorpb.CollectionVersionFile
16131611
if existingVersionFileName == "" {
16141612
// The VersionFile has not been created.

‎go/pkg/sysdb/coordinator/table_catalog_test.go‎

Lines changed: 5 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -95,6 +95,7 @@ func TestCatalog_GetCollections(t *testing.T) {
9595
name := "test_collection"
9696
testKey := "test_key"
9797
testValue := "test_value"
98+
dbId := types.NewUniqueID()
9899
collectionConfigurationJsonStr := "{\"a\": \"param\", \"b\": \"param2\", \"3\": true}"
99100
collectionAndMetadataList := []*dbmodel.CollectionAndMetadata{
100101
{
@@ -103,6 +104,7 @@ func TestCatalog_GetCollections(t *testing.T) {
103104
Name: &name,
104105
ConfigurationJsonStr: &collectionConfigurationJsonStr,
105106
Ts: types.Timestamp(1234567890),
107+
DatabaseID: dbId.String(),
106108
},
107109
CollectionMetadata: []*dbmodel.CollectionMetadata{
108110
{
@@ -136,6 +138,7 @@ func TestCatalog_GetCollections(t *testing.T) {
136138
ConfigurationJsonStr: collectionConfigurationJsonStr,
137139
Ts: types.Timestamp(1234567890),
138140
Metadata: metadata,
141+
DatabaseId: dbId,
139142
},
140143
}, collections)
141144

@@ -345,7 +348,7 @@ func TestCatalog_FlushCollectionCompactionForVersionedCollection(t *testing.T) {
345348
mockMetaDomain.On("TenantDb", mock.Anything).Return(mockTenantDb)
346349
mockMetaDomain.On("SegmentDb", mock.Anything).Return(mockSegmentDb)
347350

348-
mockCollectionDb.On("GetCollections", types.FromUniqueID(collectionID), mock.Anything, mock.Anything, mock.Anything, mock.Anything, mock.Anything).Return(mockCollectionsAndMetadata, nil)
351+
mockCollectionDb.On("GetCollectionEntries", types.FromUniqueID(collectionID), mock.Anything, mock.Anything, mock.Anything, mock.Anything, mock.Anything).Return(mockCollectionsAndMetadata, nil)
349352
mockSegmentDb.On("GetSegments", mock.Anything, mock.Anything, mock.Anything, collectionID).Return(mockSegments, nil)
350353
mockCollectionDb.On("UpdateLogPositionAndVersionInfo",
351354
collectionID.String(),
@@ -549,7 +552,7 @@ func TestCatalog_FlushCollectionCompactionForVersionedCollectionWithEmptyFilePat
549552
mockMetaDomain.On("TenantDb", mock.Anything).Return(mockTenantDb)
550553
mockMetaDomain.On("SegmentDb", mock.Anything).Return(mockSegmentDb)
551554

552-
mockCollectionDb.On("GetCollections", types.FromUniqueID(collectionID), mock.Anything, mock.Anything, mock.Anything, mock.Anything, mock.Anything).Return(mockCollectionsAndMetadata, nil)
555+
mockCollectionDb.On("GetCollectionEntries", types.FromUniqueID(collectionID), mock.Anything, mock.Anything, mock.Anything, mock.Anything, mock.Anything).Return(mockCollectionsAndMetadata, nil)
553556
mockSegmentDb.On("GetSegments", mock.Anything, mock.Anything, mock.Anything, collectionID).Return(mockSegments, nil)
554557
mockCollectionDb.On("UpdateLogPositionAndVersionInfo",
555558
collectionID.String(),

‎go/pkg/sysdb/grpc/tenant_database_service_test.go‎

Lines changed: 5 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -14,6 +14,7 @@ import (
1414
"github.com/chroma-core/chroma/go/pkg/sysdb/metastore/db/dbcore"
1515
"github.com/chroma-core/chroma/go/pkg/sysdb/metastore/db/dbmodel"
1616
s3metastore "github.com/chroma-core/chroma/go/pkg/sysdb/metastore/s3"
17+
"github.com/google/uuid"
1718
"github.com/pingcap/log"
1819
"github.com/stretchr/testify/suite"
1920
"google.golang.org/genproto/googleapis/rpc/code"
@@ -114,6 +115,8 @@ func (suite *TenantDatabaseServiceTestSuite) TestServer_TenantLastCompactionTime
114115
func (suite *TenantDatabaseServiceTestSuite) TestServer_DeleteDatabase() {
115116
tenantName := "TestDeleteDatabase"
116117
databaseName := "TestDeleteDatabase"
118+
// Generate random uuid for db id
119+
databaseeId := uuid.New().String()
117120

118121
_, err := suite.catalog.CreateTenant(context.Background(), &model.CreateTenant{
119122
Name: tenantName,
@@ -124,6 +127,8 @@ func (suite *TenantDatabaseServiceTestSuite) TestServer_DeleteDatabase() {
124127
_, err = suite.catalog.CreateDatabase(context.Background(), &model.CreateDatabase{
125128
Tenant: tenantName,
126129
Name: databaseName,
130+
ID: databaseeId,
131+
Ts: time.Now().Unix(),
127132
}, time.Now().Unix())
128133
suite.NoError(err)
129134

0 commit comments

Comments
 (0)