1// Copyright 2018 The Bazel Authors. All rights reserved. 2// 3// Licensed under the Apache License, Version 2.0 (the "License"); 4// you may not use this file except in compliance with the License. 5// You may obtain a copy of the License at 6// 7// http://www.apache.org/licenses/LICENSE-2.0 8// 9// Unless required by applicable law or agreed to in writing, software 10// distributed under the License is distributed on an "AS IS" BASIS, 11// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 12// See the License for the specific language governing permissions and 13// limitations under the License. 14 15package shard 16 17import ( 18 "archive/zip" 19 "bytes" 20 "errors" 21 "fmt" 22 "reflect" 23 "sort" 24 "strings" 25 "testing" 26) 27 28func TestSepSharder(t *testing.T) { 29 tcs := []struct { 30 name string 31 sep string 32 wantName string 33 }{ 34 { 35 name: "Hello", 36 sep: "/", 37 wantName: "Hello", 38 }, 39 { 40 name: "foo/bar/baz", 41 sep: "/", 42 wantName: "foo/bar", 43 }, 44 { 45 name: "com@[email protected]", 46 sep: "@", 47 wantName: "com@google", 48 }, 49 } 50 51 for _, tc := range tcs { 52 checkShard := func(name string, sc int) int { 53 if name != tc.wantName { 54 t.Errorf("makeSepSharder(%s).Shard(%s, 1): got name: %s wanted: %s", tc.sep, tc.name, name, tc.wantName) 55 } 56 return 0 57 } 58 59 s := MakeSepFunc(tc.sep, Func(checkShard)) 60 s(tc.name, 1) 61 } 62 63} 64 65func TestBadSharder(t *testing.T) { 66 srcZip, err := makeZip(map[string]string{"hello": "world"}) 67 if err != nil { 68 t.Fatalf("Could not make initial zip: %v", err) 69 } 70 71 for _, shardVal := range []int{-1, -244, 123} { 72 zr, err := zip.NewReader(bytes.NewReader(srcZip), int64(len(srcZip))) 73 if err != nil { 74 t.Fatalf("could not read initial zip: %v", err) 75 } 76 zws := []*zip.Writer{zip.NewWriter(&bytes.Buffer{})} 77 78 s := Func(func(name string, sc int) int { 79 return shardVal 80 }) 81 err = ZipShard(zr, zws, s) 82 if err == nil || !strings.Contains(err.Error(), "invalid shard index") { 83 t.Errorf("Returning shard value: %d gave: %v wanted an error with invalid shard index", shardVal, err) 84 } 85 } 86} 87 88func TestZipShard(t *testing.T) { 89 tcs := []struct { 90 name string 91 contents map[string]string 92 shardCount int 93 want map[int][]string 94 zipShardErr error 95 }{ 96 { 97 name: "Vanilla", 98 contents: map[string]string{ 99 "foo/hello": "world", 100 "bar/something": "stuff", 101 "blah/nothing": "here", 102 "blah/everything": "nowhere", 103 "hello/everything": "nowhere", 104 }, 105 shardCount: 5, 106 want: map[int][]string{ 107 0: {"hello/everything"}, 108 3: {"foo/hello", "bar/something"}, 109 4: {"blah/nothing", "blah/everything"}, 110 }, 111 }, 112 { 113 name: "no output shards", 114 contents: map[string]string{"something": "something"}, 115 shardCount: 0, 116 zipShardErr: errors.New("no output writers"), 117 }, 118 { 119 name: "empty input zip", 120 contents: map[string]string{}, 121 shardCount: 5, 122 want: map[int][]string{}, 123 }, 124 } 125 126 for _, tc := range tcs { 127 srcZip, err := makeZip(tc.contents) 128 if err != nil { 129 t.Errorf("%s: could not create initial zip: %v", tc.name, err) 130 } 131 zr, err := zip.NewReader(bytes.NewReader(srcZip), int64(len(srcZip))) 132 133 if err != nil { 134 t.Errorf("%s: could not read initial zip: %v", tc.name, err) 135 continue 136 } 137 bufs := make([]*bytes.Buffer, tc.shardCount) 138 zws := make([]*zip.Writer, tc.shardCount) 139 for i := range zws { 140 bufs[i] = new(bytes.Buffer) 141 zws[i] = zip.NewWriter(bufs[i]) 142 } 143 s := MakeSepFunc("/", Func(func(name string, sc int) int { 144 return len(name) % sc 145 })) 146 err = ZipShard(zr, zws, s) 147 if !reflect.DeepEqual(err, tc.zipShardErr) { 148 t.Errorf("%s: got zipshard error: %v wanted: %v", tc.name, err, tc.zipShardErr) 149 continue 150 } 151 for i, s := range bufs { 152 if err := zws[i].Close(); err != nil { 153 t.Errorf("%s: shard: %d cannot close zip writer: %v", tc.name, tc.shardCount, err) 154 continue 155 } 156 z, err := zip.NewReader(bytes.NewReader(s.Bytes()), int64(s.Len())) 157 if err != nil { 158 t.Errorf("%s: shard: %d cannot create zip reader: %v", tc.name, tc.shardCount, err) 159 continue 160 } 161 var fileNames []string 162 for _, f := range z.File { 163 fileNames = append(fileNames, f.Name) 164 } 165 sort.Strings(fileNames) 166 want, _ := tc.want[i] 167 sort.Strings(want) 168 if !reflect.DeepEqual(want, fileNames) { 169 t.Errorf("%s: shard: %d got: %s wanted: %s", tc.name, i, fileNames, want) 170 } 171 } 172 } 173} 174 175func makeZip(contents map[string]string) ([]byte, error) { 176 var zin bytes.Buffer 177 zw := zip.NewWriter(&zin) 178 for name, body := range contents { 179 f, err := zw.Create(name) 180 if err != nil { 181 return nil, fmt.Errorf("%s: could not create: %v", name, err) 182 } 183 _, err = f.Write([]byte(body)) 184 if err != nil { 185 return nil, fmt.Errorf("%s: could not write: %s due to %v: ", name, body, err) 186 } 187 188 } 189 if err := zw.Close(); err != nil { 190 return nil, err 191 } 192 return zin.Bytes(), nil 193} 194