diff --git a/coordinator/internal/logic/provertask/batch_prover_task.go b/coordinator/internal/logic/provertask/batch_prover_task.go
index f59d80c42..920d83ced 100644
--- a/coordinator/internal/logic/provertask/batch_prover_task.go
+++ b/coordinator/internal/logic/provertask/batch_prover_task.go
@@ -87,11 +87,10 @@ func (bp *BatchProverTask) Assign(ctx *gin.Context, getTaskParameter *coordinato
 		if getTaskParameter.TaskID != "" {
 			tmpBatchTask, getTaskError = bp.batchOrm.GetBatchByHash(ctx.Copy(), getTaskParameter.TaskID)
 			if getTaskError != nil {
-				log.Error("failed to get expected batch", "taskID", taskCtx.hasAssignedTask.TaskID, "err", getTaskError)
+				log.Error("failed to get expected batch", "taskID", getTaskParameter.TaskID, "err", getTaskError)
 				return nil, ErrCoordinatorInternalFailure
 			} else if tmpBatchTask == nil {
-				return nil, fmt.Errorf("Expected task (%s) is already dropped",
-					taskCtx.hasAssignedTask.TaskID)
+				return nil, fmt.Errorf("Expected task (%s) is already dropped", getTaskParameter.TaskID)
 			}
 		}
 
diff --git a/coordinator/internal/logic/provertask/bundle_prover_task.go b/coordinator/internal/logic/provertask/bundle_prover_task.go
index c0ad2a093..ab4c69d1f 100644
--- a/coordinator/internal/logic/provertask/bundle_prover_task.go
+++ b/coordinator/internal/logic/provertask/bundle_prover_task.go
@@ -85,11 +85,10 @@ func (bp *BundleProverTask) Assign(ctx *gin.Context, getTaskParameter *coordinat
 		if getTaskParameter.TaskID != "" {
 			tmpBundleTask, getTaskError = bp.bundleOrm.GetBundleByHash(ctx.Copy(), getTaskParameter.TaskID)
 			if getTaskError != nil {
-				log.Error("failed to get expected bundle", "taskID", taskCtx.hasAssignedTask.TaskID, "err", getTaskError)
+				log.Error("failed to get expected bundle", "taskID", getTaskParameter.TaskID, "err", getTaskError)
 				return nil, ErrCoordinatorInternalFailure
 			} else if tmpBundleTask == nil {
-				return nil, fmt.Errorf("Expected task (%s) is already dropped",
-					taskCtx.hasAssignedTask.TaskID)
+				return nil, fmt.Errorf("Expected task (%s) is already dropped", getTaskParameter.TaskID)
 			}
 		}
 
diff --git a/coordinator/internal/logic/provertask/chunk_prover_task.go b/coordinator/internal/logic/provertask/chunk_prover_task.go
index a914e4240..8d7ec1d99 100644
--- a/coordinator/internal/logic/provertask/chunk_prover_task.go
+++ b/coordinator/internal/logic/provertask/chunk_prover_task.go
@@ -83,11 +83,10 @@ func (cp *ChunkProverTask) Assign(ctx *gin.Context, getTaskParameter *coordinato
 		if getTaskParameter.TaskID != "" {
 			tmpChunkTask, getTaskError = cp.chunkOrm.GetChunkByHash(ctx.Copy(), getTaskParameter.TaskID)
 			if getTaskError != nil {
-				log.Error("failed to get expected chunk", "taskID", taskCtx.hasAssignedTask.TaskID, "err", getTaskError)
+				log.Error("failed to get expected chunk", "taskID", getTaskParameter.TaskID, "err", getTaskError)
 				return nil, ErrCoordinatorInternalFailure
 			} else if tmpChunkTask == nil {
-				return nil, fmt.Errorf("Expected task (%s) is already dropped",
-					taskCtx.hasAssignedTask.TaskID)
+				return nil, fmt.Errorf("Expected task (%s) is already dropped", getTaskParameter.TaskID)
 			}
 		}
 
diff --git a/crates/prover-bin/src/main.rs b/crates/prover-bin/src/main.rs
index afad98f18..8f4b0bc8d 100644
--- a/crates/prover-bin/src/main.rs
+++ b/crates/prover-bin/src/main.rs
@@ -5,7 +5,7 @@ mod zk_circuits_handler;
 use clap::{ArgAction, Parser, Subcommand};
 use prover::{LocalProver, LocalProverConfig};
 use scroll_proving_sdk::{
-    prover::ProverBuilder,
+    prover::{ProverBuilder, types::ProofType},
     utils::{get_version, init_tracing},
 };
 use std::{fs::File, path::Path, io::BufReader};
@@ -86,11 +86,11 @@ async fn main() -> eyre::Result<()> {
 
             let prover = std::sync::Arc::new(prover);
             println!("Handling task set 1: chunks ...");
-            prover.clone().one_shot(&handle_set.chunks).await;
+            assert!(prover.clone().one_shot(&handle_set.chunks, ProofType::Chunk).await);
             println!("Done! Handling task set 2: batches ...");
-            prover.clone().one_shot(&handle_set.batches).await;
+            assert!(prover.clone().one_shot(&handle_set.batches, ProofType::Batch).await);
             println!("Done! Handling task set 3: bundles ...");
-            prover.clone().one_shot(&handle_set.bundles).await;
+            assert!(prover.clone().one_shot(&handle_set.bundles, ProofType::Bundle).await);
             println!("All done!");
         }
         None => {
diff --git a/rollup/tests/integration_tool/imports.go b/rollup/tests/integration_tool/imports.go
index 178efbd7a..88cbaa221 100644
--- a/rollup/tests/integration_tool/imports.go
+++ b/rollup/tests/integration_tool/imports.go
@@ -78,7 +78,7 @@ func importData(ctx context.Context, beginBlk, endBlk uint64, chkNum, batchNum,
 	}
 
 	batchSep := randomPickKfromN(chkNum, batchNum, rng)
-	batchSep = append(batchSep, batchNum)
+	batchSep = append(batchSep, chkNum)
 	log.Info("separated batch", "border", batchSep)
 
 	headChk := int(0)
@@ -99,7 +99,7 @@ func importData(ctx context.Context, beginBlk, endBlk uint64, chkNum, batchNum,
 	}
 
 	bundleSep := randomPickKfromN(batchNum, bundleNum, rng)
-	bundleSep = append(bundleSep, bundleNum)
+	bundleSep = append(bundleSep, batchNum)
 	log.Info("separated bundle", "border", bundleSep)
 
 	headBatch := int(0)
@@ -141,6 +141,10 @@ func importChunk(ctx context.Context, db *gorm.DB, beginBlk, endBlk uint64, prev
 	if err != nil {
 		return nil, nil, err
 	}
+	err = blockOrm.UpdateChunkHashInRange(ctx, beginBlk, endBlk, dbChk.Hash)
+	if err != nil {
+		return nil, nil, err
+	}
 	log.Info("insert chunk", "From", beginBlk, "To", endBlk)
 	return dbChk, theChunk, nil
 }
@@ -173,6 +177,11 @@ func importBatch(ctx context.Context, db *gorm.DB, chks []*orm.Chunk, encChks []
 	if err != nil {
 		return nil, err
 	}
+	err = orm.NewChunk(db).UpdateBatchHashInRange(ctx, chks[0].Index, chks[len(chks)-1].Index, dbBatch.Hash)
+	if err != nil {
+		return nil, err
+	}
+
 	log.Info("insert batch", "index", index)
 	return dbBatch, nil
 }
@@ -184,6 +193,11 @@ func importBundle(ctx context.Context, db *gorm.DB, batches []*orm.Batch) (strin
 	if err != nil {
 		return "", err
 	}
+	err = orm.NewBatch(db).UpdateBundleHashInRange(ctx, batches[0].Index, batches[len(batches)-1].Index, bundle.Hash)
+	if err != nil {
+		return "", err
+	}
+
 	log.Info("insert bundle", "hash", bundle.Hash)
 	return bundle.Hash, nil
 }
diff --git a/tests/prover-e2e/.gitignore b/tests/prover-e2e/.gitignore
index c795b054e..d8e367a78 100644
--- a/tests/prover-e2e/.gitignore
+++ b/tests/prover-e2e/.gitignore
@@ -1 +1,2 @@
-build
\ No newline at end of file
+build/*
+testset.json
\ No newline at end of file