1 package throttler
2
3 import (
4 "fmt"
5 "math/rand"
6 "reflect"
7 "strconv"
8 "testing"
9 "time"
10 )
11
12 func TestThrottle(t *testing.T) {
13 var tests = []struct {
14 Desc string
15 Jobs []string
16 MaxWorkers int
17 TotalJobs int
18 }{
19 {
20 "Standard implementation",
21 []string{"job01", "job02", "job03", "job04", "job05", "job06", "job07", "job08", "job09", "job10",
22 "job11", "job12", "job13", "job14", "job15", "job16", "job17", "job18", "job19", "job20",
23 "job21", "job22", "job23", "job24", "job25", "job26", "job27", "job28", "job29", "job30",
24 "job31", "job32", "job33", "job34", "job35", "job36", "job37", "job38", "job39", "job40",
25 "job41", "job42", "job43", "job44", "job45", "job46", "job47", "job48", "job49", "job50"},
26 5,
27 -1,
28 }, {
29 "Incorrectly has 0 as TotalWorkers",
30 []string{"job01", "job02", "job03", "job04", "job05", "job06", "job07", "job08", "job09", "job10",
31 "job11", "job12", "job13", "job14", "job15", "job16", "job17", "job18", "job19", "job20",
32 "job21", "job22", "job23", "job24", "job25", "job26", "job27", "job28", "job29", "job30",
33 "job31", "job32", "job33", "job34", "job35", "job36", "job37", "job38", "job39", "job40",
34 "job41", "job42", "job43", "job44", "job45", "job46", "job47", "job48", "job49", "job50"},
35 5,
36 0,
37 }, {
38 "More workers than jobs",
39 []string{"job01", "job02", "job03", "job04", "job05", "job06", "job07", "job08", "job09", "job10",
40 "job11", "job12", "job13", "job14", "job15", "job16", "job17", "job18", "job19", "job20",
41 "job21", "job22", "job23", "job24", "job25", "job26", "job27", "job28", "job29", "job30",
42 "job31", "job32", "job33", "job34", "job35", "job36", "job37", "job38", "job39", "job40",
43 "job41", "job42", "job43", "job44", "job45", "job46", "job47", "job48", "job49", "job50"},
44 50000,
45 -1,
46 },
47 }
48
49 for _, test := range tests {
50 totalJobs := len(test.Jobs)
51 if test.TotalJobs != -1 {
52 totalJobs = test.TotalJobs
53 }
54 th := New(test.MaxWorkers, totalJobs)
55 for _, job := range test.Jobs {
56 go func(job string, th *Throttler) {
57 defer th.Done(nil)
58 time.Sleep(time.Duration(rand.Intn(5)) * time.Millisecond)
59 }(job, th)
60 th.Throttle()
61 }
62 if th.Err() != nil {
63 fmt.Println("err:", th.Err())
64 }
65 }
66 }
67
68 func TestThrottleWithErrors(t *testing.T) {
69 var tests = []struct {
70 Desc string
71 Jobs []string
72 MaxWorkers int
73 TotalJobs int
74 }{
75 {
76 "Standard implementation",
77 []string{"job01", "job02", "job03", "job04", "job05", "job06", "job07", "job08", "job09", "job10",
78 "job11", "job12", "job13", "job14", "job15", "job16", "job17", "job18", "job19", "job20",
79 "job21", "job22", "job23", "job24", "job25", "job26", "job27", "job28", "job29", "job30",
80 "job31", "job32", "job33", "job34", "job35", "job36", "job37", "job38", "job39", "job40",
81 "job41", "job42", "job43", "job44", "job45", "job46", "job47", "job48", "job49", "job50"},
82 5,
83 -1,
84 }, {
85 "Standard implementation",
86 []string{"job01", "job02"},
87 5,
88 -1,
89 },
90 }
91
92 for _, test := range tests {
93 totalJobs := len(test.Jobs)
94 if test.TotalJobs != -1 {
95 totalJobs = test.TotalJobs
96 }
97 th := New(test.MaxWorkers, totalJobs)
98 for _, job := range test.Jobs {
99 go func(job string, th *Throttler) {
100 jobNum, _ := strconv.ParseInt(job[len(job)-2:], 10, 8)
101 var err error
102 if jobNum%2 != 0 {
103 err = fmt.Errorf("Error on %s", job)
104 }
105 defer th.Done(err)
106
107 time.Sleep(time.Duration(rand.Intn(5)) * time.Millisecond)
108 }(job, th)
109 th.Throttle()
110 }
111 if len(th.Errs()) != totalJobs/2 {
112 t.Fatal("The wrong number of errors were returned")
113 }
114 if th.Err() != nil {
115 fmt.Println("err:", th.Err())
116 }
117 }
118 }
119
120 func TestThrottlePanic(t *testing.T) {
121 defer func() {
122 if r := recover(); r == nil {
123 t.Fatal("Test failed to panic")
124 }
125 }()
126 New(0, 100)
127 }
128
129 func TestBatchedThrottler(t *testing.T) {
130 var tests = []struct {
131 Desc string
132 ToBeBatched []string
133 MaxWorkers int
134 BatchSize int
135 ExpectedBatchedSlices [][]string
136 }{
137 {
138 "Standard implementation",
139 []string{"item01", "item02", "item03", "item04", "item05", "item06", "item07", "item08", "item09", "item10",
140 "item11", "item12", "item13", "item14", "item15", "item16", "item17", "item18", "item19", "item20",
141 "item21", "item22", "item23", "item24", "item25", "item26", "item27", "item28", "item29", "item30",
142 "item31", "item32", "item33", "item34", "item35", "item36", "item37", "item38", "item39", "item40",
143 "item41", "item42", "item43", "item44", "item45", "item46", "item47", "item48", "item49",
144 },
145 10,
146 2,
147 [][]string{
148 {"item01", "item02"},
149 {"item03", "item04"},
150 {"item05", "item06"},
151 {"item07", "item08"},
152 {"item09", "item10"},
153 {"item11", "item12"},
154 {"item13", "item14"},
155 {"item15", "item16"},
156 {"item17", "item18"},
157 {"item19", "item20"},
158 {"item21", "item22"},
159 {"item23", "item24"},
160 {"item25", "item26"},
161 {"item27", "item28"},
162 {"item29", "item30"},
163 {"item31", "item32"},
164 {"item33", "item34"},
165 {"item35", "item36"},
166 {"item37", "item38"},
167 {"item39", "item40"},
168 {"item41", "item42"},
169 {"item43", "item44"},
170 {"item45", "item46"},
171 {"item47", "item48"},
172 {"item49"},
173 },
174 },
175 }
176
177 for _, test := range tests {
178 th := NewBatchedThrottler(test.MaxWorkers, len(test.ToBeBatched), test.BatchSize)
179 for i := 0; i < th.TotalJobs(); i++ {
180 go func(tbbSlice []string, expectedSlice []string) {
181 var err error
182 if !reflect.DeepEqual(tbbSlice, expectedSlice) {
183 err = fmt.Errorf("wanted: %#v | got: %#v", expectedSlice, tbbSlice)
184 }
185 th.Done(err)
186 }(test.ToBeBatched[th.BatchStartIndex():th.BatchEndIndex()], test.ExpectedBatchedSlices[i])
187 if errCount := th.Throttle(); errCount > 0 {
188 break
189 }
190 }
191
192 if th.Err() != nil {
193 t.Fatal(th.Err())
194 }
195 }
196 }
197
198 func TestSetMaxWorkers(t *testing.T) {
199 var tests = []struct {
200 Desc string
201 Jobs []string
202 InitialMaxWorkers int
203 EndMaxWorkers int
204 TotalJobs int
205 }{
206 {
207 "Standard implementation",
208 []string{"job01", "job02", "job03", "job04", "job05", "job06", "job07", "job08", "job09", "job10",
209 "job11", "job12", "job13", "job14", "job15", "job16", "job17", "job18", "job19", "job20",
210 "job21", "job22", "job23", "job24", "job25", "job26", "job27", "job28", "job29", "job30",
211 "job31", "job32", "job33", "job34", "job35", "job36", "job37", "job38", "job39", "job40",
212 "job41", "job42", "job43", "job44", "job45", "job46", "job47", "job48", "job49", "job50"},
213 1,
214 5,
215 -1,
216 }, {
217 "Incorrectly has 0 as TotalWorkers",
218 []string{"job01", "job02", "job03", "job04", "job05", "job06", "job07", "job08", "job09", "job10",
219 "job11", "job12", "job13", "job14", "job15", "job16", "job17", "job18", "job19", "job20",
220 "job21", "job22", "job23", "job24", "job25", "job26", "job27", "job28", "job29", "job30",
221 "job31", "job32", "job33", "job34", "job35", "job36", "job37", "job38", "job39", "job40",
222 "job41", "job42", "job43", "job44", "job45", "job46", "job47", "job48", "job49", "job50"},
223 1,
224 5,
225 0,
226 }, {
227 "More workers than jobs",
228 []string{"job01", "job02", "job03", "job04", "job05", "job06", "job07", "job08", "job09", "job10",
229 "job11", "job12", "job13", "job14", "job15", "job16", "job17", "job18", "job19", "job20",
230 "job21", "job22", "job23", "job24", "job25", "job26", "job27", "job28", "job29", "job30",
231 "job31", "job32", "job33", "job34", "job35", "job36", "job37", "job38", "job39", "job40",
232 "job41", "job42", "job43", "job44", "job45", "job46", "job47", "job48", "job49", "job50"},
233 1,
234 50000,
235 -1,
236 },
237 }
238
239 for _, test := range tests {
240 totalJobs := len(test.Jobs)
241 if test.TotalJobs != -1 {
242 totalJobs = test.TotalJobs
243 }
244 th := New(test.InitialMaxWorkers, totalJobs)
245 for i, job := range test.Jobs {
246 if i == test.InitialMaxWorkers+1 {
247 th.SetMaxWorkers(test.EndMaxWorkers)
248 }
249 go func(job string, th *Throttler) {
250 defer th.Done(nil)
251 time.Sleep(time.Duration(rand.Intn(5)) * time.Millisecond)
252 }(job, th)
253 th.Throttle()
254 }
255 if th.Err() != nil {
256 fmt.Println("err:", th.Err())
257 }
258 }
259 }
260
261 func TestSetMaxWorkersPanic(t *testing.T) {
262 defer func() {
263 if r := recover(); r == nil {
264 t.Fatal("Test failed to panic")
265 }
266 }()
267 th := New(1, 10)
268 th.SetMaxWorkers(-1)
269 }
270
View as plain text