@@ -10,6 +10,8 @@ import (
10
10
"time"
11
11
12
12
"github.com/golang-queue/queue/core"
13
+ "github.com/golang-queue/queue/mocks"
14
+ "github.com/golang/mock/gomock"
13
15
14
16
"github.com/stretchr/testify/assert"
15
17
)
@@ -237,8 +239,6 @@ func TestHandleTimeout(t *testing.T) {
237
239
done <- w .handle (job )
238
240
}()
239
241
240
- assert .NoError (t , w .Shutdown ())
241
-
242
242
err = <- done
243
243
assert .Error (t , err )
244
244
assert .Equal (t , context .DeadlineExceeded , err )
@@ -276,8 +276,6 @@ func TestJobComplete(t *testing.T) {
276
276
done <- w .handle (job )
277
277
}()
278
278
279
- assert .NoError (t , w .Shutdown ())
280
-
281
279
err = <- done
282
280
assert .Error (t , err )
283
281
assert .Equal (t , errors .New ("job completed" ), err )
@@ -308,7 +306,7 @@ func TestTaskJobComplete(t *testing.T) {
308
306
go func () {
309
307
done <- w .handle (job )
310
308
}()
311
- assert . NoError ( t , w . Shutdown ())
309
+
312
310
err = <- done
313
311
assert .NoError (t , err )
314
312
@@ -385,3 +383,68 @@ func TestDecreaseWorkerCount(t *testing.T) {
385
383
assert .Equal (t , 2 , q .BusyWorkers ())
386
384
q .Release ()
387
385
}
386
+
387
+ func TestHandleAllJobBeforeShutdownConsumer (t * testing.T ) {
388
+ controller := gomock .NewController (t )
389
+ defer controller .Finish ()
390
+
391
+ m := mocks .NewMockQueuedMessage (controller )
392
+
393
+ w := NewConsumer (
394
+ WithFn (func (ctx context.Context , m core.QueuedMessage ) error {
395
+ time .Sleep (10 * time .Millisecond )
396
+ return nil
397
+ }),
398
+ )
399
+
400
+ done := make (chan struct {})
401
+ assert .NoError (t , w .Queue (m ))
402
+ assert .NoError (t , w .Queue (m ))
403
+ go func () {
404
+ assert .NoError (t , w .Shutdown ())
405
+ done <- struct {}{}
406
+ }()
407
+
408
+ task , err := w .Request ()
409
+ assert .NotNil (t , task )
410
+ assert .NoError (t , err )
411
+ task , err = w .Request ()
412
+ assert .NotNil (t , task )
413
+ assert .NoError (t , err )
414
+ task , err = w .Request ()
415
+ assert .Nil (t , task )
416
+ assert .True (t , errors .Is (err , ErrQueueHasBeenClosed ))
417
+ <- done
418
+ }
419
+
420
+ func TestHandleAllJobBeforeShutdownConsumerInQueue (t * testing.T ) {
421
+ controller := gomock .NewController (t )
422
+ defer controller .Finish ()
423
+
424
+ m := mocks .NewMockQueuedMessage (controller )
425
+ m .EXPECT ().Bytes ().Return ([]byte ("test" )).AnyTimes ()
426
+
427
+ messages := make (chan string , 10 )
428
+
429
+ w := NewConsumer (
430
+ WithFn (func (ctx context.Context , m core.QueuedMessage ) error {
431
+ time .Sleep (10 * time .Millisecond )
432
+ messages <- string (m .Bytes ())
433
+ return nil
434
+ }),
435
+ )
436
+
437
+ q , err := NewQueue (
438
+ WithLogger (NewLogger ()),
439
+ WithWorker (w ),
440
+ WithWorkerCount (1 ),
441
+ )
442
+ assert .NoError (t , err )
443
+
444
+ assert .NoError (t , q .Queue (m ))
445
+ assert .NoError (t , q .Queue (m ))
446
+ assert .Len (t , messages , 0 )
447
+ q .Start ()
448
+ q .Release ()
449
+ assert .Len (t , messages , 2 )
450
+ }
0 commit comments