17967: Update generated config.
[arvados.git] / services / keepstore / handler_test.go
index 76cacc569532972664996a7d4a5eede9efc33323..897447dd11c7a95a5b113d867fb0de28cbed6844 100644 (file)
@@ -21,7 +21,7 @@ import (
        "net/http"
        "net/http/httptest"
        "os"
-       "regexp"
+       "sort"
        "strings"
        "time"
 
@@ -71,10 +71,11 @@ func (s *HandlerSuite) SetUpTest(c *check.C) {
 // A RequestTester represents the parameters for an HTTP request to
 // be issued on behalf of a unit test.
 type RequestTester struct {
-       uri         string
-       apiToken    string
-       method      string
-       requestBody []byte
+       uri            string
+       apiToken       string
+       method         string
+       requestBody    []byte
+       storageClasses string
 }
 
 // Test GetBlockHandler on the following situations:
@@ -318,6 +319,54 @@ func (s *HandlerSuite) TestPutAndDeleteSkipReadonlyVolumes(c *check.C) {
        }
 }
 
+func (s *HandlerSuite) TestReadsOrderedByStorageClassPriority(c *check.C) {
+       s.cluster.Volumes = map[string]arvados.Volume{
+               "zzzzz-nyw5e-111111111111111": {
+                       Driver:         "mock",
+                       Replication:    1,
+                       StorageClasses: map[string]bool{"class1": true}},
+               "zzzzz-nyw5e-222222222222222": {
+                       Driver:         "mock",
+                       Replication:    1,
+                       StorageClasses: map[string]bool{"class2": true, "class3": true}},
+       }
+
+       for _, trial := range []struct {
+               priority1 int // priority of class1, thus vol1
+               priority2 int // priority of class2
+               priority3 int // priority of class3 (vol2 priority will be max(priority2, priority3))
+               get1      int // expected number of "get" ops on vol1
+               get2      int // expected number of "get" ops on vol2
+       }{
+               {100, 50, 50, 1, 0},   // class1 has higher priority => try vol1 first, no need to try vol2
+               {100, 100, 100, 1, 0}, // same priority, vol1 is first lexicographically => try vol1 first and succeed
+               {66, 99, 33, 1, 1},    // class2 has higher priority => try vol2 first, then try vol1
+               {66, 33, 99, 1, 1},    // class3 has highest priority => vol2 has highest => try vol2 first, then try vol1
+       } {
+               c.Logf("%+v", trial)
+               s.cluster.StorageClasses = map[string]arvados.StorageClassConfig{
+                       "class1": {Priority: trial.priority1},
+                       "class2": {Priority: trial.priority2},
+                       "class3": {Priority: trial.priority3},
+               }
+               c.Assert(s.handler.setup(context.Background(), s.cluster, "", prometheus.NewRegistry(), testServiceURL), check.IsNil)
+               IssueRequest(s.handler,
+                       &RequestTester{
+                               method:         "PUT",
+                               uri:            "/" + TestHash,
+                               requestBody:    TestBlock,
+                               storageClasses: "class1",
+                       })
+               IssueRequest(s.handler,
+                       &RequestTester{
+                               method: "GET",
+                               uri:    "/" + TestHash,
+                       })
+               c.Check(s.handler.volmgr.mountMap["zzzzz-nyw5e-111111111111111"].Volume.(*MockVolume).CallCount("Get"), check.Equals, trial.get1)
+               c.Check(s.handler.volmgr.mountMap["zzzzz-nyw5e-222222222222222"].Volume.(*MockVolume).CallCount("Get"), check.Equals, trial.get2)
+       }
+}
+
 // Test TOUCH requests.
 func (s *HandlerSuite) TestTouchHandler(c *check.C) {
        c.Assert(s.handler.setup(context.Background(), s.cluster, "", prometheus.NewRegistry(), testServiceURL), check.IsNil)
@@ -495,12 +544,8 @@ func (s *HandlerSuite) TestIndexHandler(c *check.C) {
 
        expected := `^` + TestHash + `\+\d+ \d+\n` +
                TestHash2 + `\+\d+ \d+\n\n$`
-       match, _ := regexp.MatchString(expected, response.Body.String())
-       if !match {
-               c.Errorf(
-                       "permissions on, superuser request: expected %s, got:\n%s",
-                       expected, response.Body.String())
-       }
+       c.Check(response.Body.String(), check.Matches, expected, check.Commentf(
+               "permissions on, superuser request"))
 
        // superuser /index/prefix request
        // => OK
@@ -511,12 +556,8 @@ func (s *HandlerSuite) TestIndexHandler(c *check.C) {
                response)
 
        expected = `^` + TestHash + `\+\d+ \d+\n\n$`
-       match, _ = regexp.MatchString(expected, response.Body.String())
-       if !match {
-               c.Errorf(
-                       "permissions on, superuser /index/prefix request: expected %s, got:\n%s",
-                       expected, response.Body.String())
-       }
+       c.Check(response.Body.String(), check.Matches, expected, check.Commentf(
+               "permissions on, superuser /index/prefix request"))
 
        // superuser /index/{no-such-prefix} request
        // => OK
@@ -754,25 +795,25 @@ func (s *HandlerSuite) TestPullHandler(c *check.C) {
        var testcases = []pullTest{
                {
                        "Valid pull list from an ordinary user",
-                       RequestTester{"/pull", userToken, "PUT", goodJSON},
+                       RequestTester{"/pull", userToken, "PUT", goodJSON, ""},
                        http.StatusUnauthorized,
                        "Unauthorized\n",
                },
                {
                        "Invalid pull request from an ordinary user",
-                       RequestTester{"/pull", userToken, "PUT", badJSON},
+                       RequestTester{"/pull", userToken, "PUT", badJSON, ""},
                        http.StatusUnauthorized,
                        "Unauthorized\n",
                },
                {
                        "Valid pull request from the data manager",
-                       RequestTester{"/pull", s.cluster.SystemRootToken, "PUT", goodJSON},
+                       RequestTester{"/pull", s.cluster.SystemRootToken, "PUT", goodJSON, ""},
                        http.StatusOK,
                        "Received 3 pull requests\n",
                },
                {
                        "Invalid pull request from the data manager",
-                       RequestTester{"/pull", s.cluster.SystemRootToken, "PUT", badJSON},
+                       RequestTester{"/pull", s.cluster.SystemRootToken, "PUT", badJSON, ""},
                        http.StatusBadRequest,
                        "",
                },
@@ -866,25 +907,25 @@ func (s *HandlerSuite) TestTrashHandler(c *check.C) {
        var testcases = []trashTest{
                {
                        "Valid trash list from an ordinary user",
-                       RequestTester{"/trash", userToken, "PUT", goodJSON},
+                       RequestTester{"/trash", userToken, "PUT", goodJSON, ""},
                        http.StatusUnauthorized,
                        "Unauthorized\n",
                },
                {
                        "Invalid trash list from an ordinary user",
-                       RequestTester{"/trash", userToken, "PUT", badJSON},
+                       RequestTester{"/trash", userToken, "PUT", badJSON, ""},
                        http.StatusUnauthorized,
                        "Unauthorized\n",
                },
                {
                        "Valid trash list from the data manager",
-                       RequestTester{"/trash", s.cluster.SystemRootToken, "PUT", goodJSON},
+                       RequestTester{"/trash", s.cluster.SystemRootToken, "PUT", goodJSON, ""},
                        http.StatusOK,
                        "Received 3 trash requests\n",
                },
                {
                        "Invalid trash list from the data manager",
-                       RequestTester{"/trash", s.cluster.SystemRootToken, "PUT", badJSON},
+                       RequestTester{"/trash", s.cluster.SystemRootToken, "PUT", badJSON, ""},
                        http.StatusBadRequest,
                        "",
                },
@@ -921,6 +962,9 @@ func IssueRequest(handler http.Handler, rt *RequestTester) *httptest.ResponseRec
        if rt.apiToken != "" {
                req.Header.Set("Authorization", "OAuth2 "+rt.apiToken)
        }
+       if rt.storageClasses != "" {
+               req.Header.Set("X-Keep-Storage-Classes", rt.storageClasses)
+       }
        handler.ServeHTTP(response, req)
        return response
 }
@@ -1113,6 +1157,63 @@ func (s *HandlerSuite) TestGetHandlerNoBufferLeak(c *check.C) {
        }
 }
 
+func (s *HandlerSuite) TestPutStorageClasses(c *check.C) {
+       s.cluster.Volumes = map[string]arvados.Volume{
+               "zzzzz-nyw5e-000000000000000": {Replication: 1, Driver: "mock"}, // "default" is implicit
+               "zzzzz-nyw5e-111111111111111": {Replication: 1, Driver: "mock", StorageClasses: map[string]bool{"special": true, "extra": true}},
+               "zzzzz-nyw5e-222222222222222": {Replication: 1, Driver: "mock", StorageClasses: map[string]bool{"readonly": true}, ReadOnly: true},
+       }
+       c.Assert(s.handler.setup(context.Background(), s.cluster, "", prometheus.NewRegistry(), testServiceURL), check.IsNil)
+       rt := RequestTester{
+               method:      "PUT",
+               uri:         "/" + TestHash,
+               requestBody: TestBlock,
+       }
+
+       for _, trial := range []struct {
+               ask    string
+               expect string
+       }{
+               {"", ""},
+               {"default", "default=1"},
+               {" , default , default , ", "default=1"},
+               {"special", "extra=1, special=1"},
+               {"special, readonly", "extra=1, special=1"},
+               {"special, nonexistent", "extra=1, special=1"},
+               {"extra, special", "extra=1, special=1"},
+               {"default, special", "default=1, extra=1, special=1"},
+       } {
+               c.Logf("success case %#v", trial)
+               rt.storageClasses = trial.ask
+               resp := IssueRequest(s.handler, &rt)
+               if trial.expect == "" {
+                       // any non-empty value is correct
+                       c.Check(resp.Header().Get("X-Keep-Storage-Classes-Confirmed"), check.Not(check.Equals), "")
+               } else {
+                       c.Check(sortCommaSeparated(resp.Header().Get("X-Keep-Storage-Classes-Confirmed")), check.Equals, trial.expect)
+               }
+       }
+
+       for _, trial := range []struct {
+               ask string
+       }{
+               {"doesnotexist"},
+               {"doesnotexist, readonly"},
+               {"readonly"},
+       } {
+               c.Logf("failure case %#v", trial)
+               rt.storageClasses = trial.ask
+               resp := IssueRequest(s.handler, &rt)
+               c.Check(resp.Code, check.Equals, http.StatusServiceUnavailable)
+       }
+}
+
+func sortCommaSeparated(s string) string {
+       slice := strings.Split(s, ", ")
+       sort.Strings(slice)
+       return strings.Join(slice, ", ")
+}
+
 func (s *HandlerSuite) TestPutResponseHeader(c *check.C) {
        c.Assert(s.handler.setup(context.Background(), s.cluster, "", prometheus.NewRegistry(), testServiceURL), check.IsNil)