fix: ut

bnb-chain · Jun 19, 2024 · 1abff23 · 1abff23
1 parent e757d28
commit 1abff23
Show file tree

Hide file tree

Showing 2 changed files with 29 additions and 22 deletions.
diff --git a/modular/receiver/receive_task.go b/modular/receiver/receive_task.go
@@ -114,11 +114,13 @@ func (r *ReceiveModular) HandleDoneReceivePieceTask(ctx context.Context, task ta
 		log.CtxErrorw(ctx, "failed to get checksum from db", "task", task, "error", err)
 		return nil, ErrGfSpDBWithDetail("failed to get checksum from db, error: " + err.Error())
 	}
+	// // If it already have integrity data,Avoid repetitive writing db
 	skipInsertIntegrityMeta := false
 	if len(pieceChecksums) != int(segmentCount) {
-		// If it already have data, skip this check
+		// Interface idempotent processing. If it already have integrity data, can skip this check
 		integrityMeta, integrityErr := r.baseApp.GfSpDB().GetObjectIntegrity(task.GetObjectInfo().Id.Uint64(), task.GetRedundancyIdx())
 		if integrityMeta != nil && integrityErr == nil {
+			// The checksum is obtained from integrityMeta
 			pieceChecksums = integrityMeta.PieceChecksumList
 			skipInsertIntegrityMeta = true
 		} else {
@@ -145,25 +147,27 @@ func (r *ReceiveModular) HandleDoneReceivePieceTask(ctx context.Context, task ta
 	}
 
 	setIntegrityTime := time.Now()
-	if !skipInsertIntegrityMeta && task.GetObjectInfo().GetIsUpdating() {
-		integrityMeta := &corespdb.ShadowIntegrityMeta{
-			ObjectID:          task.GetObjectInfo().Id.Uint64(),
-			RedundancyIndex:   task.GetRedundancyIdx(),
-			IntegrityChecksum: integrityChecksum,
-			PieceChecksumList: pieceChecksums,
-			Version:           task.GetObjectInfo().GetVersion(),
-			ObjectSize:        task.GetObjectInfo().GetPayloadSize(),
-		}
-		err = r.baseApp.GfSpDB().SetShadowObjectIntegrity(integrityMeta)
-	} else if !skipInsertIntegrityMeta {
-		integrityMeta := &corespdb.IntegrityMeta{
-			ObjectID:          task.GetObjectInfo().Id.Uint64(),
-			RedundancyIndex:   task.GetRedundancyIdx(),
-			IntegrityChecksum: integrityChecksum,
-			PieceChecksumList: pieceChecksums,
-			ObjectSize:        task.GetObjectInfo().GetPayloadSize(),
+	if !skipInsertIntegrityMeta {
+		if task.GetObjectInfo().GetIsUpdating() {
+			integrityMeta := &corespdb.ShadowIntegrityMeta{
+				ObjectID:          task.GetObjectInfo().Id.Uint64(),
+				RedundancyIndex:   task.GetRedundancyIdx(),
+				IntegrityChecksum: integrityChecksum,
+				PieceChecksumList: pieceChecksums,
+				Version:           task.GetObjectInfo().GetVersion(),
+				ObjectSize:        task.GetObjectInfo().GetPayloadSize(),
+			}
+			err = r.baseApp.GfSpDB().SetShadowObjectIntegrity(integrityMeta)
+		} else {
+			integrityMeta := &corespdb.IntegrityMeta{
+				ObjectID:          task.GetObjectInfo().Id.Uint64(),
+				RedundancyIndex:   task.GetRedundancyIdx(),
+				IntegrityChecksum: integrityChecksum,
+				PieceChecksumList: pieceChecksums,
+				ObjectSize:        task.GetObjectInfo().GetPayloadSize(),
+			}
+			err = r.baseApp.GfSpDB().SetObjectIntegrity(integrityMeta)
 		}
-		err = r.baseApp.GfSpDB().SetObjectIntegrity(integrityMeta)
 	}
 	metrics.PerfReceivePieceTimeHistogram.WithLabelValues("receive_piece_server_done_set_integrity_time").Observe(time.Since(setIntegrityTime).Seconds())
 	if err != nil {

diff --git a/modular/receiver/receive_task_test.go b/modular/receiver/receive_task_test.go
@@ -7,15 +7,17 @@ import (
 
 	sdkmath "cosmossdk.io/math"
 	"github.com/bnb-chain/greenfield-common/go/hash"
+	storagetypes "github.com/bnb-chain/greenfield/x/storage/types"
+	"github.com/stretchr/testify/assert"
+	"go.uber.org/mock/gomock"
+	"gorm.io/gorm"
+
 	"github.com/bnb-chain/greenfield-storage-provider/base/gfspclient"
 	"github.com/bnb-chain/greenfield-storage-provider/base/gfsppieceop"
 	"github.com/bnb-chain/greenfield-storage-provider/base/types/gfsptask"
 	"github.com/bnb-chain/greenfield-storage-provider/core/piecestore"
 	"github.com/bnb-chain/greenfield-storage-provider/core/spdb"
 	"github.com/bnb-chain/greenfield-storage-provider/core/taskqueue"
-	storagetypes "github.com/bnb-chain/greenfield/x/storage/types"
-	"github.com/stretchr/testify/assert"
-	"go.uber.org/mock/gomock"
 )
 
 func TestErrPieceStoreWithDetail(t *testing.T) {
@@ -236,6 +238,7 @@ func TestHandleDoneReceivePieceTask_PieceCountMismatch(t *testing.T) {
 	mockSPDB := spdb.NewMockSPDB(ctrl)
 	r.baseApp.SetGfSpDB(mockSPDB)
 	mockSPDB.EXPECT().GetAllReplicatePieceChecksumOptimized(gomock.Any(), gomock.Any(), gomock.Any()).Return(nil, nil).Times(1)
+	mockSPDB.EXPECT().GetObjectIntegrity(gomock.Any(), gomock.Any()).Return(nil, gorm.ErrRecordNotFound).AnyTimes()
 	_, err := r.HandleDoneReceivePieceTask(context.TODO(), mockTask)
 	assert.NotNil(t, err)
 }