X-Git-Url: https://git.arvados.org/arvados.git/blobdiff_plain/a56406d730f2a07dd442b9e99ef9dab7b7d81895..1fd5716e1714337b6ff96f6725e1f22c7a6ceb65:/services/datamanager/collection/testing.go

diff --git a/services/datamanager/collection/testing.go b/services/datamanager/collection/testing.go
index 1d3c2e811b..2238433722 100644
--- a/services/datamanager/collection/testing.go
+++ b/services/datamanager/collection/testing.go
@@ -7,6 +7,7 @@ import (
 	"git.curoverse.com/arvados.git/sdk/go/blockdigest"
 )
 
+// TestCollectionSpec with test blocks and desired replication level
 type TestCollectionSpec struct {
 	// The desired replication level
 	ReplicationLevel int
@@ -15,23 +16,23 @@ type TestCollectionSpec struct {
 	Blocks []int
 }
 
-// Creates a ReadCollections object for testing based on the give
-// specs.  Only the ReadAllCollections and UuidToCollection fields are
-// populated.  To populate other fields call rc.Summarize().
+// MakeTestReadCollections creates a ReadCollections object for testing
+// based on the give specs. Only the ReadAllCollections and UUIDToCollection
+// fields are populated. To populate other fields call rc.Summarize().
 func MakeTestReadCollections(specs []TestCollectionSpec) (rc ReadCollections) {
 	rc = ReadCollections{
 		ReadAllCollections: true,
-		UuidToCollection:   map[string]Collection{},
+		UUIDToCollection:   map[string]Collection{},
 	}
 
 	for i, spec := range specs {
 		c := Collection{
-			Uuid:              fmt.Sprintf("col%d", i),
-			OwnerUuid:         fmt.Sprintf("owner%d", i),
+			UUID:              fmt.Sprintf("col%d", i),
+			OwnerUUID:         fmt.Sprintf("owner%d", i),
 			ReplicationLevel:  spec.ReplicationLevel,
 			BlockDigestToSize: map[blockdigest.BlockDigest]int{},
 		}
-		rc.UuidToCollection[c.Uuid] = c
+		rc.UUIDToCollection[c.UUID] = c
 		for _, j := range spec.Blocks {
 			c.BlockDigestToSize[blockdigest.MakeTestBlockDigest(j)] = j
 		}
@@ -44,3 +45,17 @@ func MakeTestReadCollections(specs []TestCollectionSpec) (rc ReadCollections) {
 	}
 	return
 }
+
+// CollectionIndicesForTesting returns a slice giving the collection
+// index of each collection that was passed in to MakeTestReadCollections.
+// rc.Summarize() must be called before this method, since Summarize()
+// assigns an index to each collection.
+func (rc ReadCollections) CollectionIndicesForTesting() (indices []int) {
+	// TODO(misha): Assert that rc.Summarize() has been called.
+	numCollections := len(rc.CollectionIndexToUUID)
+	indices = make([]int, numCollections)
+	for i := 0; i < numCollections; i++ {
+		indices[i] = rc.CollectionUUIDToIndex[fmt.Sprintf("col%d", i)]
+	}
+	return
+}