fix: test

Signed-off-by: Jim Ma <[email protected]>
dragonflyoss · Dec 9, 2021 · 030a09e · 030a09e
1 parent 541ca02
commit 030a09e
Showing 1 changed file with 8 additions and 2 deletions.
diff --git a/client/daemon/peer/piece_manager_test.go b/client/daemon/peer/piece_manager_test.go
@@ -83,7 +83,7 @@ func TestPieceManager_DownloadSource(t *testing.T) {
 		checkDigest       bool
 	}{
 		{
-			name:              "multiple pieces with content length",
+			name:              "multiple pieces with content length, check digest",
 			pieceSize:         1024,
 			checkDigest:       true,
 			withContentLength: true,
@@ -95,7 +95,7 @@ func TestPieceManager_DownloadSource(t *testing.T) {
 			withContentLength: true,
 		},
 		{
-			name:              "multiple pieces without content length",
+			name:              "multiple pieces without content length, check digest",
 			pieceSize:         1024,
 			checkDigest:       true,
 			withContentLength: false,
@@ -131,12 +131,18 @@ func TestPieceManager_DownloadSource(t *testing.T) {
 		t.Run(tc.name, func(t *testing.T) {
 			/********** prepare test start **********/
 			mockPeerTask := NewMockTask(ctrl)
+			var totalPieces int32
 			mockPeerTask.EXPECT().SetContentLength(gomock.Any()).AnyTimes().DoAndReturn(
 				func(arg0 int64) error {
 					return nil
 				})
 			mockPeerTask.EXPECT().SetTotalPieces(gomock.Any()).AnyTimes().DoAndReturn(
 				func(arg0 int32) {
+					totalPieces = arg0
+				})
+			mockPeerTask.EXPECT().GetTotalPieces().AnyTimes().DoAndReturn(
+				func() int32 {
+					return totalPieces
 				})
 			mockPeerTask.EXPECT().GetPeerID().AnyTimes().DoAndReturn(
 				func() string {