| 344 | } |
| 345 | |
| 346 | func TestCancellation(t *testing.T) { |
| 347 | prePoolOpts := goleak.IgnoreCurrent() |
| 348 | p := NewPool(&Config{ |
| 349 | MaxWorkers: 1, |
| 350 | QueueDepth: 10, |
| 351 | }) |
| 352 | opts := goleak.IgnoreCurrent() |
| 353 | |
| 354 | callCount := 0 |
| 355 | cancelAfter := 2 |
| 356 | |
| 357 | ctx, cancel := context.WithCancel(context.Background()) |
| 358 | |
| 359 | ret := []byte{0x01, 0x02} |
| 360 | fn := func(_ context.Context, _ interface{}) (interface{}, error) { |
| 361 | callCount++ |
| 362 | if callCount >= cancelAfter { |
| 363 | cancel() |
| 364 | } |
| 365 | return ret, nil |
| 366 | } |
| 367 | payloads := []interface{}{1, 2, 3, 4, 5, 6, 7} |
| 368 | |
| 369 | results, funcErrs, err := p.RunJobs(ctx, payloads, fn) |
| 370 | require.Len(t, results, 2) |
| 371 | for i := range results { |
| 372 | assert.Equal(t, ret, results[i]) |
| 373 | } |
| 374 | |
| 375 | assert.Nil(t, err) |
| 376 | assert.Nil(t, funcErrs) |
| 377 | goleak.VerifyNone(t, opts) |
| 378 | |
| 379 | p.Shutdown() |
| 380 | goleak.VerifyNone(t, prePoolOpts) |
| 381 | } |