xref: /aosp_15_r20/external/bazelbuild-rules_android/src/tools/ak/bucketize/partitioner_test.go (revision 9e965d6fece27a77de5377433c2f7e6999b8cc0b)
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 bucketize
16
17import (
18	"archive/zip"
19	"bytes"
20	"encoding/xml"
21	"fmt"
22	"io"
23	"io/ioutil"
24	"reflect"
25	"sort"
26	"strconv"
27	"strings"
28	"testing"
29
30	"src/common/golang/shard"
31	"src/tools/ak/res/res"
32)
33
34func TestInternalStorePathResource(t *testing.T) {
35	// test internal storePathResource and skip the creation of real files.
36	tcs := []struct {
37		name       string
38		inFiles    map[string]string
39		partitions map[res.Type][]io.Writer
40		shardFn    shard.Func
41		want       map[res.Type][][]string
42		wantErr    bool
43	}{
44		{
45			name: "MultipleResTypeFilesWithShardsOfDifferentSizes",
46			inFiles: map[string]string{
47				"res/drawable/2-foo.xml":  "all",
48				"res/layout/0-bar.xml":    "your",
49				"res/color/0-baz.xml":     "base",
50				"res/layout/1-qux.xml":    "are",
51				"res/drawable/0-quux.xml": "belong",
52				"res/color/0-corge.xml":   "to",
53				"res/color/0-grault.xml":  "us",
54				"res/layout/0-garply.xml": "!",
55			},
56			shardFn: shard.Func(func(fqn string, shardCount int) int {
57				// sharding strategy is built into the file name as "<shard num>-foo.bar" (i.e. 8-baz.xml)
58				name := strings.Split(fqn, "/")[1]
59				ai := strings.SplitN(name, "-", 2)[0]
60				shard, err := strconv.Atoi(ai)
61				if err != nil {
62					t.Fatalf("Atoi(%s) got err: %v", ai, err)
63				}
64				return shard
65			}),
66			partitions: map[res.Type][]io.Writer{
67				res.Drawable: {&bytes.Buffer{}, &bytes.Buffer{}, &bytes.Buffer{}},
68				res.Color:    {&bytes.Buffer{}},
69				res.Layout:   {&bytes.Buffer{}, &bytes.Buffer{}},
70			},
71			want: map[res.Type][][]string{
72				res.Drawable: {{"res/drawable/0-quux.xml"}, {}, {"res/drawable/2-foo.xml"}},
73				res.Color:    {{"res/color/0-baz.xml", "res/color/0-corge.xml", "res/color/0-grault.xml"}},
74				res.Layout:   {{"res/layout/0-bar.xml", "res/layout/0-garply.xml"}, {"res/layout/1-qux.xml"}},
75			},
76		},
77		{
78			name: "IgnoredFilePatterns",
79			inFiles: map[string]string{
80				"res/drawable/.ignore": "me",
81			},
82			shardFn:    shard.FNV,
83			partitions: map[res.Type][]io.Writer{res.Drawable: {&bytes.Buffer{}}},
84			wantErr:    true,
85		},
86		{
87			name:       "NoFiles",
88			inFiles:    map[string]string{},
89			shardFn:    shard.FNV,
90			partitions: map[res.Type][]io.Writer{res.Drawable: {&bytes.Buffer{}}},
91			want:       map[res.Type][][]string{res.Drawable: {{}}},
92		},
93	}
94
95	order := make(map[string]int)
96	for _, tc := range tcs {
97		t.Run(tc.name, func(t *testing.T) {
98			ps, err := makePartitionSession(tc.partitions, tc.shardFn, order)
99			if err != nil {
100				t.Errorf("MakePartitionSession(%v, %v, %d) got err: %v", tc.partitions, tc.shardFn, 0, err)
101				return
102			}
103
104			for k, v := range tc.inFiles {
105				pi, err := res.ParsePath(k)
106				if err != nil {
107					if !tc.wantErr {
108						t.Fatalf("ParsePath(%s) got err: %v", k, err)
109					}
110					return
111				}
112				if err := ps.storePathResource(pi, strings.NewReader(v)); err != nil {
113					t.Fatalf("storePathResource got unexpected err: %v", err)
114				}
115			}
116
117			if err := ps.Close(); err != nil {
118				t.Errorf("partition Close() got err: %v", err)
119				return
120			}
121
122			// validate data outputted to the partitions
123			got := make(map[res.Type][][]string)
124			for rt, shards := range tc.partitions {
125				shardPaths := make([][]string, 0, len(shards))
126				for _, shard := range shards {
127					br := bytes.NewReader(shard.(*bytes.Buffer).Bytes())
128					rr, err := zip.NewReader(br, br.Size())
129					if err != nil {
130						t.Errorf("NewReader(%v, %d) got err: %v", br, br.Size(), err)
131						return
132					}
133					paths := make([]string, 0, len(rr.File))
134					for _, f := range rr.File {
135						paths = append(paths, f.Name)
136						c, err := readAll(f)
137						if err != nil {
138							t.Errorf("readAll got err: %v", err)
139							return
140						}
141						if tc.inFiles[f.Name] != c {
142							t.Errorf("error copying data for %s got %q but wanted %q", f.Name, c, tc.inFiles[f.Name])
143							return
144						}
145					}
146					sort.Strings(paths)
147					shardPaths = append(shardPaths, paths)
148				}
149				got[rt] = shardPaths
150			}
151			if !reflect.DeepEqual(got, tc.want) {
152				t.Errorf("DeepEqual(\n%#v\n,\n%#v\n): returned false", got, tc.want)
153			}
154		})
155	}
156}
157
158func TestCollectValues(t *testing.T) {
159	tcs := []struct {
160		name       string
161		pathVPsMap map[string]map[res.FullyQualifiedName][]byte
162		pathRAMap  map[string][]xml.Attr
163		partitions map[res.Type][]io.Writer
164		want       map[res.Type][][]string
165		wantErr    bool
166	}{
167		{
168			name: "MultipleResTypesShardsResources",
169			partitions: map[res.Type][]io.Writer{
170				res.Attr:   {&bytes.Buffer{}, &bytes.Buffer{}},
171				res.String: {&bytes.Buffer{}, &bytes.Buffer{}},
172				res.Color:  {&bytes.Buffer{}, &bytes.Buffer{}},
173			},
174			pathVPsMap: map[string]map[res.FullyQualifiedName][]byte{
175				"res/values/strings.xml": {
176					res.FullyQualifiedName{Package: "res-auto", Type: res.String, Name: "foo"}: []byte("<string name='foo'>bar</string>"),
177					res.FullyQualifiedName{Package: "android", Type: res.String, Name: "baz"}:  []byte("<string name='baz'>qux</string>"),
178					res.FullyQualifiedName{Package: "res-auto", Type: res.Attr, Name: "quux"}:  []byte("<attr name='quux'>corge</attr>"),
179				},
180				"res/values/attr.xml": {
181					res.FullyQualifiedName{Package: "android", Type: res.Attr, Name: "foo"}: []byte("<attr name='android:foo'>bar</attr>"),
182				},
183				"baz/res/values/attr.xml": {
184					res.FullyQualifiedName{Package: "android", Type: res.Attr, Name: "bazfoo"}: []byte("<attr name='android:bazfoo'>qix</attr>"),
185				},
186				"baz/res/values/strings.xml": {
187					res.FullyQualifiedName{Package: "android", Type: res.String, Name: "baz"}: []byte("<string name='baz'>qux</string>"),
188				},
189				"foo/res/values/attr.xml": {
190					res.FullyQualifiedName{Package: "android", Type: res.Attr, Name: "foofoo"}: []byte("<attr name='android:foofoo'>qex</attr>"),
191				},
192				"foo/res/values/color.xml": {
193					res.FullyQualifiedName{Package: "android", Type: res.Color, Name: "foobar"}: []byte("<color name='foobar'>#FFFFFFFF</color>"),
194				},
195				"dir/res/values/strings.xml": {
196					res.FullyQualifiedName{Package: "android", Type: res.String, Name: "dirbaz"}: []byte("<string name='dirbaz'>qux</string>"),
197				},
198				"dir/res/values/color.xml": {
199					res.FullyQualifiedName{Package: "android", Type: res.Color, Name: "dirfoobar"}: []byte("<color name='dirfoobar'>#FFFFFFFF</color>"),
200				},
201			},
202			pathRAMap: map[string][]xml.Attr{
203				"res/values/strings.xml": {
204					xml.Attr{Name: xml.Name{Space: "xmlns", Local: "ns1"}, Value: "path1"},
205					xml.Attr{Name: xml.Name{Space: "xmlns", Local: "ns2"}, Value: "path2"},
206				},
207			},
208			want: map[res.Type][][]string{
209				res.Attr: {
210					{
211						"res/values/strings.xml", "<?xml version='1.0' encoding='utf-8'?><resources xmlns:ns1=\"path1\" xmlns:ns2=\"path2\"><attr name='quux'>corge</attr></resources>",
212					},
213					{
214						"res/values/strings.xml", "<?xml version='1.0' encoding='utf-8'?><resources xmlns:ns1=\"path1\" xmlns:ns2=\"path2\"></resources>",
215					},
216				},
217				res.String: {
218					{
219						"res/values/strings.xml", "<?xml version='1.0' encoding='utf-8'?><resources xmlns:ns1=\"path1\" xmlns:ns2=\"path2\"><string name='baz'>qux</string><string name='foo'>bar</string></resources>",
220						"res/values/strings.xml", "<?xml version='1.0' encoding='utf-8'?><resources><string name='dirbaz'>qux</string></resources>",
221						"res/values/strings.xml", "<?xml version='1.0' encoding='utf-8'?><resources><string name='baz'>qux</string></resources>",
222					},
223					{
224						"res/values/strings.xml", "<?xml version='1.0' encoding='utf-8'?><resources xmlns:ns1=\"path1\" xmlns:ns2=\"path2\"></resources>",
225						"res/values/strings.xml", "<?xml version='1.0' encoding='utf-8'?><resources></resources>",
226						"res/values/strings.xml", "<?xml version='1.0' encoding='utf-8'?><resources></resources>",
227					},
228				},
229				res.Color: {
230					{
231						"res/values/color.xml", "<?xml version='1.0' encoding='utf-8'?><resources><color name='foobar'>#FFFFFFFF</color></resources>",
232						"res/values/color.xml", "<?xml version='1.0' encoding='utf-8'?><resources><color name='dirfoobar'>#FFFFFFFF</color></resources>",
233					},
234					{
235						"res/values/color.xml", "<?xml version='1.0' encoding='utf-8'?><resources></resources>",
236						"res/values/color.xml", "<?xml version='1.0' encoding='utf-8'?><resources></resources>",
237					},
238				},
239			},
240		},
241		{
242			name: "NoValuesPayloads",
243			pathVPsMap: map[string]map[res.FullyQualifiedName][]byte{
244				"res/values/strings.xml": {},
245			},
246			partitions: map[res.Type][]io.Writer{res.String: {&bytes.Buffer{}}},
247			want:       map[res.Type][][]string{res.String: {{}}},
248		},
249		{
250			name: "ResTypeValuesResTypeMismatch",
251			pathVPsMap: map[string]map[res.FullyQualifiedName][]byte{
252				"res/values/strings.xml": {
253					res.FullyQualifiedName{
254						Package: "res-auto",
255						Type:    res.String,
256						Name:    "foo",
257					}: []byte("<string name='foo'>bar</string>"),
258				},
259			},
260			partitions: map[res.Type][]io.Writer{res.Attr: {&bytes.Buffer{}}},
261			want:       map[res.Type][][]string{res.Attr: {{}}},
262		},
263	}
264
265	shardFn := func(name string, shardCount int) int { return 0 }
266	order := map[string]int{
267		"foo/res/values/attr.xml":    0,
268		"foo/res/values/color.xml":   1,
269		"res/values/attr.xml":        2,
270		"res/values/strings.xml":     3,
271		"dir/res":                    4,
272		"baz/res/values/attr.xml":    5,
273		"baz/res/values/strings.xml": 6,
274	}
275	for _, tc := range tcs {
276		t.Run(tc.name, func(t *testing.T) {
277			ps, err := makePartitionSession(tc.partitions, shardFn, order)
278			if err != nil {
279				t.Errorf("makePartitionSession(%v, %v, %d) got err: %v", tc.partitions, shard.FNV, 0, err)
280				return
281			}
282			for p, vps := range tc.pathVPsMap {
283				pi, err := res.ParsePath(p)
284				if err != nil {
285					t.Errorf("ParsePath(%s) got err: %v", p, err)
286					return
287				}
288				for fqn, p := range vps {
289					ps.CollectValues(&res.ValuesResource{Src: &pi, N: fqn, Payload: p})
290				}
291			}
292			for p, as := range tc.pathRAMap {
293				pi, err := res.ParsePath(p)
294				if err != nil {
295					t.Errorf("ParsePath(%s) got err: %v", p, err)
296					return
297				}
298				for _, a := range as {
299					ps.CollectResourcesAttribute(&ResourcesAttribute{ResFile: &pi, Attribute: a})
300				}
301			}
302			if err := ps.Close(); err != nil {
303				t.Errorf("partition Close() got err: %v", err)
304				return
305			}
306
307			// validate data outputted to the partitions.
308			got := make(map[res.Type][][]string)
309			for rt, shards := range tc.partitions {
310				shardPaths := make([][]string, 0, len(shards))
311				for _, shard := range shards {
312					br := bytes.NewReader(shard.(*bytes.Buffer).Bytes())
313					rr, err := zip.NewReader(br, br.Size())
314					if err != nil {
315						t.Errorf("NewReader(%v, %d) got err: %v", br, br.Size(), err)
316						return
317					}
318					paths := make([]string, 0, len(rr.File))
319					for _, f := range rr.File {
320						c, err := readAll(f)
321						if err != nil {
322							t.Errorf("readAll got err: %v", err)
323							return
324						}
325						paths = append(paths, f.Name, c)
326					}
327					shardPaths = append(shardPaths, paths)
328				}
329				got[rt] = shardPaths
330			}
331			if !reflect.DeepEqual(got, tc.want) {
332				t.Errorf("DeepEqual(\n%#v\n,\n%#v\n): returned false", got, tc.want)
333			}
334		})
335	}
336}
337
338func readAll(f *zip.File) (string, error) {
339	rc, err := f.Open()
340	if err != nil {
341		return "", fmt.Errorf("%q: Open got err: %v", f.Name, err)
342	}
343	defer rc.Close()
344	body, err := ioutil.ReadAll(rc)
345	if err != nil {
346		return "", fmt.Errorf("%q: ReadAll got err: %v", f.Name, err)
347	}
348	return string(body), nil
349}
350