Try fix test

chroma-core · Dec 4, 2024 · 8de0623 · 8de0623
1 parent 7e5b8cc
commit 8de0623
Show file tree

Hide file tree

Showing 2 changed files with 12 additions and 4 deletions.
diff --git a/go/pkg/sysdb/coordinator/coordinator_test.go b/go/pkg/sysdb/coordinator/coordinator_test.go
@@ -310,10 +310,15 @@ func (suite *APIsTestSuite) TestCreateCollectionAndSegments() {
 	suite.NoError(error)
 	suite.Equal(newCollection.ID, collection.ID)
 	suite.Equal(newCollection.Name, collection.Name)
-	suite.Len(collection_segments, 2)
-	for i, segment := range collection_segments {
-		suite.Equal(segments[i].ID, segment.ID)
+	expected_ids, actual_ids := []types.UniqueID{}, []types.UniqueID{}
+	for _, segment := range segments {
+		expected_ids = append(expected_ids, segment.ID)
+	}
+	for _, segment := range collection_segments {
+		suite.Equal(collection.ID, segment.CollectionID)
+		actual_ids = append(expected_ids, segment.ID)
 	}
+	suite.ElementsMatch(expected_ids, actual_ids)
 
 	// Attempt to create a duplicate collection (should fail)
 	_, _, err = suite.coordinator.CreateCollectionAndSegments(ctx, newCollection, segments)

diff --git a/go/pkg/sysdb/coordinator/table_catalog.go b/go/pkg/sysdb/coordinator/table_catalog.go
@@ -328,7 +328,10 @@ func (tc *Catalog) GetCollectionWithSegments(ctx context.Context, collectionID t
 		if e != nil {
 			return e
 		}
-		if len(collections) != 1 {
+		if len(collections) == 0 {
+			return common.ErrCollectionNotFound
+		}
+		if len(collections) > 1 {
 			return common.ErrCollectionUniqueConstraintViolation
 		}
 		collection = collections[0]