Merge branch 'pr/28'
[arvados.git] / services / datamanager / collection / testing.go
1 // Code used for testing only.
2
3 package collection
4
5 import (
6         "fmt"
7         "git.curoverse.com/arvados.git/sdk/go/blockdigest"
8 )
9
10 // TestCollectionSpec with test blocks and desired replication level
11 type TestCollectionSpec struct {
12         // The desired replication level
13         ReplicationLevel int
14         // Blocks this contains, represented by ints. Ints repeated will
15         // still only represent one block
16         Blocks []int
17 }
18
19 // MakeTestReadCollections creates a ReadCollections object for testing
20 // based on the give specs. Only the ReadAllCollections and UUIDToCollection
21 // fields are populated. To populate other fields call rc.Summarize().
22 func MakeTestReadCollections(specs []TestCollectionSpec) (rc ReadCollections) {
23         rc = ReadCollections{
24                 ReadAllCollections: true,
25                 UUIDToCollection:   map[string]Collection{},
26         }
27
28         for i, spec := range specs {
29                 c := Collection{
30                         UUID:              fmt.Sprintf("col%d", i),
31                         OwnerUUID:         fmt.Sprintf("owner%d", i),
32                         ReplicationLevel:  spec.ReplicationLevel,
33                         BlockDigestToSize: map[blockdigest.BlockDigest]int{},
34                 }
35                 rc.UUIDToCollection[c.UUID] = c
36                 for _, j := range spec.Blocks {
37                         c.BlockDigestToSize[blockdigest.MakeTestBlockDigest(j)] = j
38                 }
39                 // We compute the size in a separate loop because the value
40                 // computed in the above loop would be invalid if c.Blocks
41                 // contained duplicates.
42                 for _, size := range c.BlockDigestToSize {
43                         c.TotalSize += size
44                 }
45         }
46         return
47 }
48
49 // CollectionIndicesForTesting returns a slice giving the collection
50 // index of each collection that was passed in to MakeTestReadCollections.
51 // rc.Summarize() must be called before this method, since Summarize()
52 // assigns an index to each collection.
53 func (rc ReadCollections) CollectionIndicesForTesting() (indices []int) {
54         // TODO(misha): Assert that rc.Summarize() has been called.
55         numCollections := len(rc.CollectionIndexToUUID)
56         indices = make([]int, numCollections)
57         for i := 0; i < numCollections; i++ {
58                 indices[i] = rc.CollectionUUIDToIndex[fmt.Sprintf("col%d", i)]
59         }
60         return
61 }