refactoring TestUpdateBarPriority
Vladimir Bauer
2 years ago
| 203 | 203 | } |
| 204 | 204 | } |
| 205 | 205 | |
| 206 | func makeUpdateBarPriorityTest(refresh, lazy bool) func(*testing.T) { | |
| 207 | return func(t *testing.T) { | |
| 208 | shutdown := make(chan interface{}) | |
| 209 | refreshCh := make(chan interface{}) | |
| 210 | ctx, cancel := context.WithCancel(context.Background()) | |
| 211 | p := mpb.NewWithContext(ctx, | |
| 212 | mpb.WithOutput(io.Discard), | |
| 213 | mpb.WithManualRefresh(refreshCh), | |
| 214 | mpb.WithShutdownNotifier(shutdown), | |
| 215 | ) | |
| 216 | a := p.AddBar(100, mpb.BarPriority(1), mpb.BarID(1)) | |
| 217 | b := p.AddBar(100, mpb.BarPriority(2), mpb.BarID(2)) | |
| 218 | c := p.AddBar(100, mpb.BarPriority(3), mpb.BarID(3)) | |
| 219 | ||
| 220 | p.UpdateBarPriority(c, 2, lazy) | |
| 221 | p.UpdateBarPriority(b, 3, lazy) | |
| 222 | checkOrder := []*mpb.Bar{b, c, a} // updated order | |
| 223 | ||
| 224 | if !refresh { | |
| 225 | if lazy { | |
| 226 | checkOrder = []*mpb.Bar{c, b, a} // pristine order | |
| 227 | } | |
| 228 | } else { | |
| 229 | refreshCh <- time.Now() | |
| 230 | } | |
| 231 | ||
| 232 | go cancel() | |
| 233 | ||
| 234 | bars := (<-shutdown).([]*mpb.Bar) | |
| 235 | if l := len(bars); l != 3 { | |
| 236 | t.Errorf("Expected len of bars: %d, got: %d", 3, l) | |
| 237 | } | |
| 238 | ||
| 239 | p.Wait() | |
| 240 | pq := mpb.PriorityQueue(bars) | |
| 241 | ||
| 242 | for _, b := range checkOrder { | |
| 243 | // higher priority pops first | |
| 244 | if bar := heap.Pop(&pq).(*mpb.Bar); bar.ID() != b.ID() { | |
| 245 | t.Errorf("Expected bar id: %d, got bar id: %d", b.ID(), bar.ID()) | |
| 246 | } | |
| 247 | } | |
| 248 | } | |
| 249 | } | |
| 250 | ||
| 206 | 251 | func TestUpdateBarPriority(t *testing.T) { |
| 207 | shutdown := make(chan interface{}) | |
| 208 | ctx, cancel := context.WithCancel(context.Background()) | |
| 209 | p := mpb.NewWithContext(ctx, | |
| 210 | mpb.WithOutput(io.Discard), // auto refresh is disabled because of io.Discard | |
| 211 | mpb.WithShutdownNotifier(shutdown), | |
| 212 | ) | |
| 213 | a := p.AddBar(100, mpb.BarPriority(1)) | |
| 214 | b := p.AddBar(100, mpb.BarPriority(2)) | |
| 215 | c := p.AddBar(100, mpb.BarPriority(3)) | |
| 216 | ||
| 217 | identity := map[*mpb.Bar]string{ | |
| 218 | a: "a", | |
| 219 | b: "b", | |
| 220 | c: "c", | |
| 221 | } | |
| 222 | ||
| 223 | p.UpdateBarPriority(c, 2, false) | |
| 224 | p.UpdateBarPriority(b, 3, false) | |
| 225 | ||
| 226 | go cancel() | |
| 227 | ||
| 228 | bars := (<-shutdown).([]*mpb.Bar) | |
| 229 | if l := len(bars); l != 3 { | |
| 230 | t.Errorf("Expected len of bars: %d, got: %d", 3, l) | |
| 231 | } | |
| 232 | ||
| 233 | p.Wait() | |
| 234 | pq := mpb.PriorityQueue(bars) | |
| 235 | ||
| 236 | if bar := heap.Pop(&pq).(*mpb.Bar); bar != b { | |
| 237 | t.Errorf("Expected bar b, got: %s", identity[bar]) | |
| 238 | } | |
| 239 | if bar := heap.Pop(&pq).(*mpb.Bar); bar != c { | |
| 240 | t.Errorf("Expected bar c, got: %s", identity[bar]) | |
| 241 | } | |
| 242 | if bar := heap.Pop(&pq).(*mpb.Bar); bar != a { | |
| 243 | t.Errorf("Expected bar a, got: %s", identity[bar]) | |
| 244 | } | |
| 252 | makeUpdateBarPriorityTest(false, false)(t) | |
| 253 | makeUpdateBarPriorityTest(true, false)(t) | |
| 245 | 254 | } |
| 246 | 255 | |
| 247 | 256 | func TestUpdateBarPriorityLazy(t *testing.T) { |
| 248 | shutdown := make(chan interface{}) | |
| 249 | refreshCh := make(chan interface{}) | |
| 250 | ctx, cancel := context.WithCancel(context.Background()) | |
| 251 | p := mpb.NewWithContext(ctx, | |
| 252 | mpb.WithOutput(io.Discard), | |
| 253 | mpb.WithManualRefresh(refreshCh), | |
| 254 | mpb.WithShutdownNotifier(shutdown), | |
| 255 | ) | |
| 256 | a := p.AddBar(100, mpb.BarPriority(1)) | |
| 257 | b := p.AddBar(100, mpb.BarPriority(2)) | |
| 258 | c := p.AddBar(100, mpb.BarPriority(3)) | |
| 259 | ||
| 260 | identity := map[*mpb.Bar]string{ | |
| 261 | a: "a", | |
| 262 | b: "b", | |
| 263 | c: "c", | |
| 264 | } | |
| 265 | ||
| 266 | // UpdateBarPriority with lazy flag needs at least one refresh cycle to | |
| 267 | // update order inside underlying binary heap | |
| 268 | p.UpdateBarPriority(c, 2, true) | |
| 269 | p.UpdateBarPriority(b, 3, true) | |
| 270 | ||
| 271 | // ommiting following line will fail the test which proves correct lazy behavior | |
| 272 | refreshCh <- time.Now() | |
| 273 | ||
| 274 | go cancel() | |
| 275 | ||
| 276 | bars := (<-shutdown).([]*mpb.Bar) | |
| 277 | if l := len(bars); l != 3 { | |
| 278 | t.Errorf("Expected len of bars: %d, got: %d", 3, l) | |
| 279 | } | |
| 280 | ||
| 281 | p.Wait() | |
| 282 | pq := mpb.PriorityQueue(bars) | |
| 283 | ||
| 284 | if bar := heap.Pop(&pq).(*mpb.Bar); bar != b { | |
| 285 | t.Errorf("Expected bar b, got: %s", identity[bar]) | |
| 286 | } | |
| 287 | if bar := heap.Pop(&pq).(*mpb.Bar); bar != c { | |
| 288 | t.Errorf("Expected bar c, got: %s", identity[bar]) | |
| 289 | } | |
| 290 | if bar := heap.Pop(&pq).(*mpb.Bar); bar != a { | |
| 291 | t.Errorf("Expected bar a, got: %s", identity[bar]) | |
| 292 | } | |
| 257 | makeUpdateBarPriorityTest(false, true)(t) | |
| 258 | makeUpdateBarPriorityTest(true, true)(t) | |
| 293 | 259 | } |
| 294 | 260 | |
| 295 | 261 | func TestNoOutput(t *testing.T) { |