diff --git a/integration_tests/pipelined_memdb_test.go b/integration_tests/pipelined_memdb_test.go
index e2838588a..f2cd3f2a4 100644
--- a/integration_tests/pipelined_memdb_test.go
+++ b/integration_tests/pipelined_memdb_test.go
@@ -353,11 +353,10 @@ func (s *testPipelinedMemDBSuite) TestPipelinedPrefetch() {
 	s.Nil(txn.GetMemBuffer().FlushWait())
 	m, err = txn.BatchGet(context.Background(), [][]byte{[]byte("99")})
 	s.Nil(err)
-	// restore this check after tikv return pairs for buffer batch get
-	//s.Equal(m, map[string][]byte{})
-	//v, err = panicWhenReadingRemoteBuffer([]byte("99"))
-	//s.Nil(err)
-	//s.Equal(v, []byte{})
+	s.Equal(m, map[string][]byte{})
+	v, err = panicWhenReadingRemoteBuffer([]byte("99"))
+	s.Nil(err)
+	s.Equal(v, []byte{})
 	txn.Rollback()
 
 	// empty memdb should also cache the not exist result.
diff --git a/internal/unionstore/pipelined_memdb_test.go b/internal/unionstore/pipelined_memdb_test.go
index 769f77d53..cf85aa52b 100644
--- a/internal/unionstore/pipelined_memdb_test.go
+++ b/internal/unionstore/pipelined_memdb_test.go
@@ -323,6 +323,7 @@ func TestPipelinedAdjustFlushCondition(t *testing.T) {
 	flushed, err = memdb.Flush(false)
 	require.Nil(t, err)
 	require.True(t, flushed)
+	require.Nil(t, memdb.FlushWait())
 
 	// need 2 keys to flush
 	require.Nil(t, failpoint.Enable("tikvclient/pipelinedMemDBMinFlushKeys", `return(2)`))
@@ -332,6 +333,7 @@ func TestPipelinedAdjustFlushCondition(t *testing.T) {
 	flushed, err = memdb.Flush(false)
 	require.Nil(t, err)
 	require.False(t, flushed)
+	require.Nil(t, memdb.FlushWait())
 
 	// need 2 keys to flush, but force threshold reached
 	require.Nil(t, failpoint.Enable("tikvclient/pipelinedMemDBMinFlushKeys", `return(2)`))