Merge branch 'master' into 1885-keep-proxy refs #1885
[arvados.git] / sdk / go / src / arvados.org / keepclient / keepclient_test.go
index 348b913fd584f20880378f2660018cf8db8e73ae..600d7393c519190ca849a4ad33fd5594cdec55c7 100644 (file)
@@ -1,7 +1,7 @@
 package keepclient
 
 import (
-       "arvados.org/buffer"
+       "arvados.org/streamer"
        "crypto/md5"
        "flag"
        "fmt"
@@ -111,17 +111,15 @@ func (this StubPutHandler) ServeHTTP(resp http.ResponseWriter, req *http.Request
 }
 
 func RunBogusKeepServer(st http.Handler, port int) (listener net.Listener, url string) {
-       server := http.Server{Handler: st}
-
        var err error
        listener, err = net.ListenTCP("tcp", &net.TCPAddr{Port: port})
        if err != nil {
                panic(fmt.Sprintf("Could not listen on tcp port %v", port))
        }
 
-       url = fmt.Sprintf("http://localhost:%d", listener.Addr().(*net.TCPAddr).Port)
+       url = fmt.Sprintf("http://localhost:%d", port)
 
-       go server.Serve(listener)
+       go http.Serve(listener, st)
        return listener, url
 }
 
@@ -141,6 +139,8 @@ func UploadToStubHelper(c *C, st http.Handler, f func(KeepClient, string,
 }
 
 func (s *StandaloneSuite) TestUploadToStubKeepServer(c *C) {
+       log.Printf("TestUploadToStubKeepServer")
+
        st := StubPutHandler{
                c,
                "acbd18db4cc2f85cedef654fccc4a4d8",
@@ -159,11 +159,15 @@ func (s *StandaloneSuite) TestUploadToStubKeepServer(c *C) {
 
                        <-st.handled
                        status := <-upload_status
-                       c.Check(status, DeepEquals, uploadStatus{nil, fmt.Sprintf("%s/%s", url, st.expectPath), 200})
+                       c.Check(status, DeepEquals, uploadStatus{nil, fmt.Sprintf("%s/%s", url, st.expectPath), 200, 1})
                })
+
+       log.Printf("TestUploadToStubKeepServer done")
 }
 
 func (s *StandaloneSuite) TestUploadToStubKeepServerBufferReader(c *C) {
+       log.Printf("TestUploadToStubKeepServerBufferReader")
+
        st := StubPutHandler{
                c,
                "acbd18db4cc2f85cedef654fccc4a4d8",
@@ -175,22 +179,23 @@ func (s *StandaloneSuite) TestUploadToStubKeepServerBufferReader(c *C) {
                func(kc KeepClient, url string, reader io.ReadCloser,
                        writer io.WriteCloser, upload_status chan uploadStatus) {
 
-                       tr := buffer.StartTransferFromReader(512, reader)
+                       tr := streamer.AsyncStreamFromReader(512, reader)
                        defer tr.Close()
 
-                       br1 := tr.MakeBufferReader()
+                       br1 := tr.MakeStreamReader()
 
                        go kc.uploadToKeepServer(url, st.expectPath, br1, upload_status, 3)
 
                        writer.Write([]byte("foo"))
                        writer.Close()
 
-                       <-tr.Reader_status
                        <-st.handled
 
                        status := <-upload_status
-                       c.Check(status, DeepEquals, uploadStatus{nil, fmt.Sprintf("%s/%s", url, st.expectPath), 200})
+                       c.Check(status, DeepEquals, uploadStatus{nil, fmt.Sprintf("%s/%s", url, st.expectPath), 200, 1})
                })
+
+       log.Printf("TestUploadToStubKeepServerBufferReader done")
 }
 
 type FailHandler struct {
@@ -203,6 +208,8 @@ func (this FailHandler) ServeHTTP(resp http.ResponseWriter, req *http.Request) {
 }
 
 func (s *StandaloneSuite) TestFailedUploadToStubKeepServer(c *C) {
+       log.Printf("TestFailedUploadToStubKeepServer")
+
        st := FailHandler{
                make(chan string)}
 
@@ -220,10 +227,10 @@ func (s *StandaloneSuite) TestFailedUploadToStubKeepServer(c *C) {
                        <-st.handled
 
                        status := <-upload_status
-                       c.Check(status.Url, Equals, fmt.Sprintf("%s/%s", url, hash))
-                       c.Check(status.StatusCode, Equals, 500)
+                       c.Check(status.url, Equals, fmt.Sprintf("%s/%s", url, hash))
+                       c.Check(status.statusCode, Equals, 500)
                })
-
+       log.Printf("TestFailedUploadToStubKeepServer done")
 }
 
 type KeepServer struct {
@@ -279,6 +286,8 @@ func (s *StandaloneSuite) TestPutB(c *C) {
                (s1 == shuff[1] && s2 == shuff[0]),
                Equals,
                true)
+
+       log.Printf("TestPutB done")
 }
 
 func (s *StandaloneSuite) TestPutHR(c *C) {
@@ -327,6 +336,8 @@ func (s *StandaloneSuite) TestPutHR(c *C) {
                (s1 == shuff[1] && s2 == shuff[0]),
                Equals,
                true)
+
+       log.Printf("TestPutHR done")
 }
 
 func (s *StandaloneSuite) TestPutWithFail(c *C) {
@@ -419,6 +430,8 @@ func (s *StandaloneSuite) TestPutWithTooManyFail(c *C) {
        c.Check(err, Equals, InsufficientReplicasError)
        c.Check(replicas, Equals, 1)
        c.Check(<-st.handled, Equals, shuff[1])
+
+       log.Printf("TestPutWithTooManyFail done")
 }
 
 type StubGetHandler struct {
@@ -436,6 +449,7 @@ func (this StubGetHandler) ServeHTTP(resp http.ResponseWriter, req *http.Request
 }
 
 func (s *StandaloneSuite) TestGet(c *C) {
+       log.Printf("TestGet")
 
        hash := fmt.Sprintf("%x", md5.Sum([]byte("foo")))
 
@@ -453,6 +467,7 @@ func (s *StandaloneSuite) TestGet(c *C) {
        kc.Service_roots = []string{url}
 
        r, n, url2, err := kc.Get(hash)
+       defer r.Close()
        c.Check(err, Equals, nil)
        c.Check(n, Equals, int64(3))
        c.Check(url2, Equals, fmt.Sprintf("%s/%s", url, hash))
@@ -460,6 +475,8 @@ func (s *StandaloneSuite) TestGet(c *C) {
        content, err2 := ioutil.ReadAll(r)
        c.Check(err2, Equals, nil)
        c.Check(content, DeepEquals, []byte("foo"))
+
+       log.Printf("TestGet done")
 }
 
 func (s *StandaloneSuite) TestGetFail(c *C) {
@@ -568,11 +585,19 @@ func (s *ServerRequiredSuite) TestPutGetHead(c *C) {
        kc, err := MakeKeepClient()
        c.Assert(err, Equals, nil)
 
-       hash, replicas, err := kc.PutB([]byte("foo"))
-       c.Check(hash, Equals, fmt.Sprintf("%x", md5.Sum([]byte("foo"))))
-       c.Check(replicas, Equals, 2)
-       c.Check(err, Equals, nil)
+       hash := fmt.Sprintf("%x", md5.Sum([]byte("foo")))
 
+       {
+               n, _, err := kc.Ask(hash)
+               c.Check(err, Equals, BlockNotFound)
+               c.Check(n, Equals, int64(0))
+       }
+       {
+               hash2, replicas, err := kc.PutB([]byte("foo"))
+               c.Check(hash2, Equals, hash)
+               c.Check(replicas, Equals, 2)
+               c.Check(err, Equals, nil)
+       }
        {
                r, n, url2, err := kc.Get(hash)
                c.Check(err, Equals, nil)
@@ -583,7 +608,6 @@ func (s *ServerRequiredSuite) TestPutGetHead(c *C) {
                c.Check(err2, Equals, nil)
                c.Check(content, DeepEquals, []byte("foo"))
        }
-
        {
                n, url2, err := kc.Ask(hash)
                c.Check(err, Equals, nil)
@@ -591,3 +615,68 @@ func (s *ServerRequiredSuite) TestPutGetHead(c *C) {
                c.Check(url2, Equals, fmt.Sprintf("http://localhost:25108/%s", hash))
        }
 }
+
+type StubProxyHandler struct {
+       handled chan string
+}
+
+func (this StubProxyHandler) ServeHTTP(resp http.ResponseWriter, req *http.Request) {
+       resp.Header().Set("X-Keep-Replicas-Stored", "2")
+       this.handled <- fmt.Sprintf("http://%s", req.Host)
+}
+
+func (s *StandaloneSuite) TestPutProxy(c *C) {
+       log.Printf("TestPutProxy")
+
+       st := StubProxyHandler{make(chan string, 1)}
+
+       kc, _ := MakeKeepClient()
+
+       kc.Want_replicas = 2
+       kc.Using_proxy = true
+       kc.ApiToken = "abc123"
+       kc.Service_roots = make([]string, 1)
+
+       ks1 := RunSomeFakeKeepServers(st, 1, 2990)
+
+       for i, k := range ks1 {
+               kc.Service_roots[i] = k.url
+               defer k.listener.Close()
+       }
+
+       _, replicas, err := kc.PutB([]byte("foo"))
+       <-st.handled
+
+       c.Check(err, Equals, nil)
+       c.Check(replicas, Equals, 2)
+
+       log.Printf("TestPutProxy done")
+}
+
+func (s *StandaloneSuite) TestPutProxyInsufficientReplicas(c *C) {
+       log.Printf("TestPutProxy")
+
+       st := StubProxyHandler{make(chan string, 1)}
+
+       kc, _ := MakeKeepClient()
+
+       kc.Want_replicas = 3
+       kc.Using_proxy = true
+       kc.ApiToken = "abc123"
+       kc.Service_roots = make([]string, 1)
+
+       ks1 := RunSomeFakeKeepServers(st, 1, 2990)
+
+       for i, k := range ks1 {
+               kc.Service_roots[i] = k.url
+               defer k.listener.Close()
+       }
+
+       _, replicas, err := kc.PutB([]byte("foo"))
+       <-st.handled
+
+       c.Check(err, Equals, InsufficientReplicasError)
+       c.Check(replicas, Equals, 2)
+
+       log.Printf("TestPutProxy done")
+}