+func (s *StandaloneSuite) TestUploadWithStorageClasses(c *C) {
+ for _, trial := range []struct {
+ respHeader string
+ expectMap map[string]int
+ }{
+ {"", nil},
+ {"foo=1", map[string]int{"foo": 1}},
+ {" foo=1 , bar=2 ", map[string]int{"foo": 1, "bar": 2}},
+ {" =foo=1 ", nil},
+ {"foo", nil},
+ } {
+ st := &StubPutHandler{
+ c: c,
+ expectPath: "acbd18db4cc2f85cedef654fccc4a4d8",
+ expectAPIToken: "abc123",
+ expectBody: "foo",
+ expectStorageClass: "",
+ returnStorageClasses: trial.respHeader,
+ handled: make(chan string),
+ }
+
+ UploadToStubHelper(c, st,
+ func(kc *KeepClient, url string, reader io.ReadCloser, writer io.WriteCloser, uploadStatusChan chan uploadStatus) {
+ go kc.uploadToKeepServer(url, st.expectPath, nil, reader, uploadStatusChan, len("foo"), kc.getRequestID())
+
+ writer.Write([]byte("foo"))
+ writer.Close()
+
+ <-st.handled
+ status := <-uploadStatusChan
+ c.Check(status, DeepEquals, uploadStatus{nil, fmt.Sprintf("%s/%s", url, st.expectPath), 200, 1, trial.expectMap, ""})
+ })
+ }
+}
+
+func (s *StandaloneSuite) TestPutWithStorageClasses(c *C) {
+ nServers := 5
+ for _, trial := range []struct {
+ replicas int
+ clientClasses []string
+ putClasses []string // putClasses takes precedence over clientClasses
+ minRequests int
+ maxRequests int
+ success bool
+ }{
+ {1, []string{"class1"}, nil, 1, 1, true},
+ {2, []string{"class1"}, nil, 1, 2, true},
+ {3, []string{"class1"}, nil, 2, 3, true},
+ {1, []string{"class1", "class2"}, nil, 1, 1, true},
+ {3, nil, []string{"class1"}, 2, 3, true},
+ {1, nil, []string{"class1", "class2"}, 1, 1, true},
+ {1, []string{"class404"}, []string{"class1", "class2"}, 1, 1, true},
+ {1, []string{"class1"}, []string{"class404", "class2"}, nServers, nServers, false},
+ {nServers*2 + 1, []string{"class1"}, nil, nServers, nServers, false},
+ {1, []string{"class404"}, nil, nServers, nServers, false},
+ {1, []string{"class1", "class404"}, nil, nServers, nServers, false},
+ {1, nil, []string{"class1", "class404"}, nServers, nServers, false},
+ } {
+ c.Logf("%+v", trial)
+ st := &StubPutHandler{
+ c: c,
+ expectPath: "acbd18db4cc2f85cedef654fccc4a4d8",
+ expectAPIToken: "abc123",
+ expectBody: "foo",
+ expectStorageClass: "*",
+ returnStorageClasses: "class1=2, class2=2",
+ handled: make(chan string, 100),
+ }
+ ks := RunSomeFakeKeepServers(st, nServers)
+ arv, _ := arvadosclient.MakeArvadosClient()
+ kc, _ := MakeKeepClient(arv)
+ kc.Want_replicas = trial.replicas
+ kc.StorageClasses = trial.clientClasses
+ arv.ApiToken = "abc123"
+ localRoots := make(map[string]string)
+ writableLocalRoots := make(map[string]string)
+ for i, k := range ks {
+ localRoots[fmt.Sprintf("zzzzz-bi6l4-fakefakefake%03d", i)] = k.url
+ writableLocalRoots[fmt.Sprintf("zzzzz-bi6l4-fakefakefake%03d", i)] = k.url
+ defer k.listener.Close()
+ }
+ kc.SetServiceRoots(localRoots, writableLocalRoots, nil)
+
+ _, err := kc.BlockWrite(context.Background(), arvados.BlockWriteOptions{
+ Data: []byte("foo"),
+ StorageClasses: trial.putClasses,
+ })
+ if trial.success {
+ c.Check(err, check.IsNil)
+ } else {
+ c.Check(err, check.NotNil)
+ }
+ c.Check(len(st.handled) >= trial.minRequests, check.Equals, true, check.Commentf("len(st.handled)==%d, trial.minRequests==%d", len(st.handled), trial.minRequests))
+ c.Check(len(st.handled) <= trial.maxRequests, check.Equals, true, check.Commentf("len(st.handled)==%d, trial.maxRequests==%d", len(st.handled), trial.maxRequests))
+ if !trial.success && trial.replicas == 1 && c.Check(len(st.requests) >= 2, check.Equals, true) {
+ // Max concurrency should be 1. First request
+ // should have succeeded for class1. Second
+ // request should only ask for class404.
+ c.Check(st.requests[1].Header.Get("X-Keep-Storage-Classes"), check.Equals, "class404")
+ }
+ }
+}
+