fix: daemon test

Signed-off-by: Jim Ma <[email protected]>
dragonflyoss · Jan 15, 2022 · e82bd4e · e82bd4e
1 parent 6c790a8
commit e82bd4e
Showing 1 changed file with 38 additions and 27 deletions.
diff --git a/client/daemon/peer/peertask_manager_test.go b/client/daemon/peer/peertask_manager_test.go
@@ -379,30 +379,6 @@ func TestPeerTaskManager_TaskSuite(t *testing.T) {
 				return sourceClient
 			},
 		},
-		{
-			name:               "normal size scope - schedule timeout - auto back source",
-			taskData:           testBytes,
-			pieceParallelCount: 4,
-			pieceSize:          1024,
-			peerID:             "peer-0",
-			peerPacketDelay:    []time.Duration{time.Second},
-			scheduleTimeout:    time.Nanosecond,
-			urlGenerator: func(ts *testSpec) string {
-				server := httptest.NewServer(http.HandlerFunc(
-					func(w http.ResponseWriter, r *http.Request) {
-						n, err := w.Write(testBytes)
-						assert.Nil(err)
-						assert.Equal(len(ts.taskData), n)
-					}))
-				ts.cleanUp = append(ts.cleanUp, func() {
-					server.Close()
-				})
-				return server.URL
-			},
-			sizeScope:            base.SizeScope_NORMAL,
-			mockPieceDownloader:  nil,
-			mockHTTPSourceClient: nil,
-		},
 		{
 			name:                "normal size scope - back source - no content length",
 			taskData:            testBytes,
@@ -449,6 +425,30 @@ func TestPeerTaskManager_TaskSuite(t *testing.T) {
 				return sourceClient
 			},
 		},
+		{
+			name:               "normal size scope - schedule timeout - auto back source",
+			taskData:           testBytes,
+			pieceParallelCount: 4,
+			pieceSize:          1024,
+			peerID:             "peer-0",
+			peerPacketDelay:    []time.Duration{time.Second},
+			scheduleTimeout:    time.Nanosecond,
+			urlGenerator: func(ts *testSpec) string {
+				server := httptest.NewServer(http.HandlerFunc(
+					func(w http.ResponseWriter, r *http.Request) {
+						n, err := w.Write(testBytes)
+						assert.Nil(err)
+						assert.Equal(len(ts.taskData), n)
+					}))
+				ts.cleanUp = append(ts.cleanUp, func() {
+					server.Close()
+				})
+				return server.URL
+			},
+			sizeScope:            base.SizeScope_NORMAL,
+			mockPieceDownloader:  nil,
+			mockHTTPSourceClient: nil,
+		},
 	}
 
 	for _, _tc := range testCases {
@@ -651,12 +651,23 @@ func (ts *testSpec) runConductorTest(assert *testifyassert.Assertions, require *
 		assert.True(r, fmt.Sprintf("task %d result should be true", i))
 	}
 
-	var taskCount int
+	var (
+		runningTaskCount int
+		success          bool
+	)
+	select {
+	case <-ptc.successCh:
+		success = true
+	case <-ptc.failCh:
+	case <-time.After(10 * time.Minute):
+	}
+	assert.True(success, "task should success")
+
 	ptm.runningPeerTasks.Range(func(key, value interface{}) bool {
-		taskCount++
+		runningTaskCount++
 		return true
 	})
-	assert.Equal(0, taskCount, "no running tasks")
+	assert.Equal(0, runningTaskCount, "no running tasks")
 
 	// test reuse stream task
 	rc, _, ok := ptm.tryReuseStreamPeerTask(context.Background(), request)