1// Copyright 2019 Google Inc. 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 main
16
17import (
18	"fmt"
19	"reflect"
20	"testing"
21)
22
23func TestSplitList(t *testing.T) {
24	testcases := []struct {
25		inputCount int
26		shardCount int
27		want       [][]string
28	}{
29		{
30			inputCount: 1,
31			shardCount: 1,
32			want:       [][]string{{"1"}},
33		},
34		{
35			inputCount: 1,
36			shardCount: 2,
37			want:       [][]string{{"1"}, {}},
38		},
39		{
40			inputCount: 4,
41			shardCount: 2,
42			want:       [][]string{{"1", "2"}, {"3", "4"}},
43		},
44		{
45			inputCount: 19,
46			shardCount: 10,
47			want: [][]string{
48				{"1", "2"},
49				{"3", "4"},
50				{"5", "6"},
51				{"7", "8"},
52				{"9", "10"},
53				{"11", "12"},
54				{"13", "14"},
55				{"15", "16"},
56				{"17", "18"},
57				{"19"},
58			},
59		},
60		{
61			inputCount: 15,
62			shardCount: 10,
63			want: [][]string{
64				{"1", "2"},
65				{"3", "4"},
66				{"5", "6"},
67				{"7", "8"},
68				{"9", "10"},
69				{"11"},
70				{"12"},
71				{"13"},
72				{"14"},
73				{"15"},
74			},
75		},
76	}
77
78	for _, tc := range testcases {
79		t.Run(fmt.Sprintf("%d/%d", tc.inputCount, tc.shardCount), func(t *testing.T) {
80			input := []string{}
81			for i := 1; i <= tc.inputCount; i++ {
82				input = append(input, fmt.Sprintf("%d", i))
83			}
84
85			got := splitList(input, tc.shardCount)
86
87			if !reflect.DeepEqual(got, tc.want) {
88				t.Errorf("unexpected result for splitList([]string{...%d...}, %d):\nwant: %v\n got: %v\n",
89					tc.inputCount, tc.shardCount, tc.want, got)
90			}
91		})
92	}
93}
94