main_test.go 1.9 KB

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