1 files changed, 4 insertions, 14 deletions

diff --git a/services/pull/check_test.go b/services/pull/check_test.go
index 52209b4d35..57ccb20033 100644
--- a/services/pull/check_test.go
+++ b/services/pull/check_test.go
@@ -5,6 +5,7 @@
 package pull
 
 import (
+	"context"
 	"strconv"
 	"testing"
 	"time"
@@ -31,7 +32,7 @@ func TestPullRequest_AddToTaskQueue(t *testing.T) {
 
 	cfg, err := setting.GetQueueSettings(setting.CfgProvider, "pr_patch_checker")
 	assert.NoError(t, err)
-	prPatchCheckerQueue, err = queue.NewWorkerPoolQueueBySetting("pr_patch_checker", cfg, testHandler, true)
+	prPatchCheckerQueue, err = queue.NewWorkerPoolQueueWithContext(context.Background(), "pr_patch_checker", cfg, testHandler, true)
 	assert.NoError(t, err)
 
 	pr := unittest.AssertExistsAndLoadBean(t, &issues_model.PullRequest{ID: 2})
@@ -46,12 +47,7 @@ func TestPullRequest_AddToTaskQueue(t *testing.T) {
 	assert.True(t, has)
 	assert.NoError(t, err)
 
-	var queueShutdown, queueTerminate []func()
-	go prPatchCheckerQueue.Run(func(shutdown func()) {
-		queueShutdown = append(queueShutdown, shutdown)
-	}, func(terminate func()) {
-		queueTerminate = append(queueTerminate, terminate)
-	})
+	go prPatchCheckerQueue.Run()
 
 	select {
 	case id := <-idChan:
@@ -67,12 +63,6 @@ func TestPullRequest_AddToTaskQueue(t *testing.T) {
 	pr = unittest.AssertExistsAndLoadBean(t, &issues_model.PullRequest{ID: 2})
 	assert.Equal(t, issues_model.PullRequestStatusChecking, pr.Status)
 
-	for _, callback := range queueShutdown {
-		callback()
-	}
-	for _, callback := range queueTerminate {
-		callback()
-	}
-
+	prPatchCheckerQueue.ShutdownWait(5 * time.Second)
 	prPatchCheckerQueue = nil
 }