Merge pull request #1 from arvados/master
[arvados.git] / sdk / go / keepclient / collectionreader_test.go
1 // Copyright (C) The Arvados Authors. All rights reserved.
2 //
3 // SPDX-License-Identifier: Apache-2.0
4
5 package keepclient
6
7 import (
8         "crypto/md5"
9         "fmt"
10         "io"
11         "io/ioutil"
12         "math/rand"
13         "net/http"
14         "os"
15         "strconv"
16         "strings"
17
18         "git.arvados.org/arvados.git/sdk/go/arvadosclient"
19         "git.arvados.org/arvados.git/sdk/go/arvadostest"
20         check "gopkg.in/check.v1"
21 )
22
23 var _ = check.Suite(&CollectionReaderUnit{})
24
25 type CollectionReaderUnit struct {
26         arv     *arvadosclient.ArvadosClient
27         kc      *KeepClient
28         handler SuccessHandler
29 }
30
31 func (s *CollectionReaderUnit) SetUpTest(c *check.C) {
32         var err error
33         s.arv, err = arvadosclient.MakeArvadosClient()
34         c.Assert(err, check.IsNil)
35         s.arv.ApiToken = arvadostest.ActiveToken
36
37         s.kc, err = MakeKeepClient(s.arv)
38         c.Assert(err, check.IsNil)
39
40         s.handler = SuccessHandler{
41                 disk: make(map[string][]byte),
42                 lock: make(chan struct{}, 1),
43                 ops:  new(int),
44         }
45         localRoots := make(map[string]string)
46         for i, k := range RunSomeFakeKeepServers(s.handler, 4) {
47                 localRoots[fmt.Sprintf("zzzzz-bi6l4-fakefakefake%03d", i)] = k.url
48         }
49         s.kc.SetServiceRoots(localRoots, localRoots, nil)
50 }
51
52 type SuccessHandler struct {
53         disk map[string][]byte
54         lock chan struct{} // channel with buffer==1: full when an operation is in progress.
55         ops  *int          // number of operations completed
56 }
57
58 func (h SuccessHandler) ServeHTTP(resp http.ResponseWriter, req *http.Request) {
59         switch req.Method {
60         case "PUT":
61                 buf, err := ioutil.ReadAll(req.Body)
62                 if err != nil {
63                         resp.WriteHeader(500)
64                         return
65                 }
66                 pdh := fmt.Sprintf("%x+%d", md5.Sum(buf), len(buf))
67                 h.lock <- struct{}{}
68                 h.disk[pdh] = buf
69                 if h.ops != nil {
70                         (*h.ops)++
71                 }
72                 <-h.lock
73                 resp.Write([]byte(pdh))
74         case "GET":
75                 pdh := req.URL.Path[1:]
76                 h.lock <- struct{}{}
77                 buf, ok := h.disk[pdh]
78                 if h.ops != nil {
79                         (*h.ops)++
80                 }
81                 <-h.lock
82                 if !ok {
83                         resp.WriteHeader(http.StatusNotFound)
84                 } else {
85                         resp.Header().Set("Content-Length", fmt.Sprintf("%d", len(buf)))
86                         resp.Write(buf)
87                 }
88         default:
89                 resp.WriteHeader(http.StatusMethodNotAllowed)
90         }
91 }
92
93 type rdrTest struct {
94         mt   string      // manifest text
95         f    string      // filename
96         want interface{} // error or string to expect
97 }
98
99 func (s *CollectionReaderUnit) TestCollectionReaderContent(c *check.C) {
100         s.kc.PutB([]byte("foo"))
101         s.kc.PutB([]byte("bar"))
102         s.kc.PutB([]byte("Hello world\n"))
103         s.kc.PutB([]byte(""))
104
105         mt := arvadostest.PathologicalManifest
106
107         for _, testCase := range []rdrTest{
108                 {mt: mt, f: "zzzz", want: os.ErrNotExist},
109                 {mt: mt, f: "frob", want: os.ErrNotExist},
110                 {mt: mt, f: "/segmented/frob", want: "frob"},
111                 {mt: mt, f: "./segmented/frob", want: "frob"},
112                 {mt: mt, f: "/f", want: "f"},
113                 {mt: mt, f: "./f", want: "f"},
114                 {mt: mt, f: "foo bar//baz", want: "foo"},
115                 {mt: mt, f: "foo/zero", want: ""},
116                 {mt: mt, f: "zero@0", want: ""},
117                 {mt: mt, f: "zero@1", want: ""},
118                 {mt: mt, f: "zero@4", want: ""},
119                 {mt: mt, f: "zero@9", want: ""},
120                 {mt: mt, f: "f", want: "f"},
121                 {mt: mt, f: "ooba", want: "ooba"},
122                 {mt: mt, f: "overlapReverse/o", want: "o"},
123                 {mt: mt, f: "overlapReverse/oo", want: "oo"},
124                 {mt: mt, f: "overlapReverse/ofoo", want: "ofoo"},
125                 {mt: mt, f: "foo bar/baz", want: "foo"},
126                 {mt: mt, f: "segmented/frob", want: "frob"},
127                 {mt: mt, f: "segmented/oof", want: "oof"},
128         } {
129                 c.Logf("%#v", testCase)
130                 rdr, err := s.kc.CollectionFileReader(map[string]interface{}{"manifest_text": testCase.mt}, testCase.f)
131                 switch want := testCase.want.(type) {
132                 case error:
133                         c.Check(rdr, check.IsNil)
134                         c.Check(err, check.Equals, want)
135                 case string:
136                         buf := make([]byte, len(want))
137                         n, err := io.ReadFull(rdr, buf)
138                         c.Check(err, check.IsNil)
139                         for i := 0; i < 4; i++ {
140                                 c.Check(string(buf), check.Equals, want)
141                                 n, err = rdr.Read(buf)
142                                 c.Check(n, check.Equals, 0)
143                                 c.Check(err, check.Equals, io.EOF)
144                         }
145
146                         for a := len(want) - 2; a >= 0; a-- {
147                                 for b := a + 1; b <= len(want); b++ {
148                                         offset, err := rdr.Seek(int64(a), io.SeekStart)
149                                         c.Logf("...a=%d, b=%d", a, b)
150                                         c.Check(err, check.IsNil)
151                                         c.Check(offset, check.Equals, int64(a))
152                                         buf := make([]byte, b-a)
153                                         n, err := io.ReadFull(rdr, buf)
154                                         c.Check(err, check.IsNil)
155                                         c.Check(n, check.Equals, b-a)
156                                         c.Check(string(buf), check.Equals, want[a:b])
157                                 }
158                         }
159                         offset, err := rdr.Seek(-1, io.SeekStart)
160                         c.Check(err, check.NotNil)
161                         c.Check(offset, check.Equals, int64(len(want)))
162
163                         c.Check(rdr.Close(), check.Equals, nil)
164                 }
165         }
166 }
167
168 func (s *CollectionReaderUnit) TestCollectionReaderManyBlocks(c *check.C) {
169         h := md5.New()
170         buf := make([]byte, 4096)
171         locs := make([]string, len(buf))
172         testdata := make([]byte, 0, len(buf)*len(buf))
173         filesize := 0
174         for i := range locs {
175                 _, err := rand.Read(buf[:i])
176                 h.Write(buf[:i])
177                 locs[i], _, err = s.kc.PutB(buf[:i])
178                 c.Assert(err, check.IsNil)
179                 filesize += i
180                 testdata = append(testdata, buf[:i]...)
181         }
182         manifest := "./random " + strings.Join(locs, " ") + " 0:" + strconv.Itoa(filesize) + ":bytes.bin\n"
183         dataMD5 := h.Sum(nil)
184
185         checkMD5 := md5.New()
186         rdr, err := s.kc.CollectionFileReader(map[string]interface{}{"manifest_text": manifest}, "random/bytes.bin")
187         c.Assert(err, check.IsNil)
188         defer rdr.Close()
189
190         _, err = io.Copy(checkMD5, rdr)
191         c.Check(err, check.IsNil)
192         _, err = rdr.Read(make([]byte, 1))
193         c.Check(err, check.Equals, io.EOF)
194         c.Check(checkMD5.Sum(nil), check.DeepEquals, dataMD5)
195
196         size, err := rdr.Seek(0, io.SeekEnd)
197         c.Check(err, check.IsNil)
198         buf = make([]byte, len(testdata))
199         copy(buf, testdata)
200         curPos := size
201         for i := 0; i < 16; i++ {
202                 offset := rand.Intn(len(buf) - 1)
203                 count := rand.Intn(len(buf) - offset)
204                 if rand.Intn(2) == 0 {
205                         curPos, err = rdr.Seek(int64(offset)-curPos, io.SeekCurrent)
206                 } else {
207                         curPos, err = rdr.Seek(int64(offset), io.SeekStart)
208                 }
209                 c.Check(curPos, check.Equals, int64(offset))
210                 for count > 0 {
211                         n, err := rdr.Read(buf[offset : offset+count])
212                         c.Assert(err, check.IsNil)
213                         c.Assert(n > 0, check.Equals, true)
214                         offset += n
215                         count -= n
216                 }
217                 curPos, err = rdr.Seek(0, io.SeekCurrent)
218                 c.Check(curPos, check.Equals, int64(offset))
219         }
220         c.Check(md5.Sum(buf), check.DeepEquals, md5.Sum(testdata))
221         c.Check(buf[:1000], check.DeepEquals, testdata[:1000])
222
223         expectPos := curPos + size + 12345
224         curPos, err = rdr.Seek(size+12345, io.SeekCurrent)
225         c.Check(err, check.IsNil)
226         c.Check(curPos, check.Equals, expectPos)
227
228         curPos, err = rdr.Seek(8-curPos, io.SeekCurrent)
229         c.Check(err, check.IsNil)
230         c.Check(curPos, check.Equals, int64(8))
231
232         curPos, err = rdr.Seek(-9, io.SeekCurrent)
233         c.Check(err, check.NotNil)
234         c.Check(curPos, check.Equals, int64(8))
235 }
236
237 func (s *CollectionReaderUnit) TestCollectionReaderCloseEarly(c *check.C) {
238         s.kc.BlockCache = &BlockCache{}
239         s.kc.PutB([]byte("foo"))
240         s.kc.PutB([]byte("bar"))
241         s.kc.PutB([]byte("baz"))
242
243         mt := ". "
244         for i := 0; i < 300; i++ {
245                 mt += "acbd18db4cc2f85cedef654fccc4a4d8+3 37b51d194a7513e45b56f6524f2d51f2+3 73feffa4b7f6bb68e44cf984c85f6e88+3 "
246         }
247         mt += "0:2700:foo900.txt\n"
248
249         // Grab the stub server's lock, ensuring our cfReader doesn't
250         // get anything back from its first call to kc.Get() before we
251         // have a chance to call Close().
252         s.handler.lock <- struct{}{}
253         opsBeforeRead := *s.handler.ops
254
255         rdr, err := s.kc.CollectionFileReader(map[string]interface{}{"manifest_text": mt}, "foo900.txt")
256         c.Assert(err, check.IsNil)
257
258         firstReadDone := make(chan struct{})
259         go func() {
260                 n, err := rdr.Read(make([]byte, 3))
261                 c.Check(n, check.Equals, 3)
262                 c.Check(err, check.IsNil)
263                 close(firstReadDone)
264         }()
265
266         // Release the stub server's lock. The first GET operation will proceed.
267         <-s.handler.lock
268
269         // Make sure our first read operation consumes the data
270         // received from the first GET.
271         <-firstReadDone
272
273         err = rdr.Close()
274         c.Check(err, check.IsNil)
275
276         // Stub should have handled exactly one GET request.
277         c.Check(*s.handler.ops, check.Equals, opsBeforeRead+1)
278 }
279
280 func (s *CollectionReaderUnit) TestCollectionReaderDataError(c *check.C) {
281         manifest := ". ffffffffffffffffffffffffffffffff+1 0:1:notfound.txt\n"
282         buf := make([]byte, 1)
283         rdr, err := s.kc.CollectionFileReader(map[string]interface{}{"manifest_text": manifest}, "notfound.txt")
284         c.Check(err, check.IsNil)
285         for i := 0; i < 2; i++ {
286                 _, err = io.ReadFull(rdr, buf)
287                 c.Check(err, check.NotNil)
288                 c.Check(err, check.Not(check.Equals), io.EOF)
289         }
290         c.Check(rdr.Close(), check.IsNil)
291 }