From 764611ba192ca454f34912067588ff694707a6de Mon Sep 17 00:00:00 2001
From: rishabh_mittal <mittalrishabh@gmail.com>
Date: Wed, 8 Jan 2025 09:13:07 -0800
Subject: [PATCH] fix test cases

---
 internal/locate/region_request3_test.go | 18 +++++++++++++++++-
 1 file changed, 17 insertions(+), 1 deletion(-)

diff --git a/internal/locate/region_request3_test.go b/internal/locate/region_request3_test.go
index 6b5f17502d..142b688781 100644
--- a/internal/locate/region_request3_test.go
+++ b/internal/locate/region_request3_test.go
@@ -1404,6 +1404,20 @@ func (s *testRegionRequestToThreeStoresSuite) TestReplicaReadFallbackToLeaderReg
 	s.NotNil(regionLoc)
 
 	s.regionRequestSender.client = &fnClient{fn: func(ctx context.Context, addr string, req *tikvrpc.Request, timeout time.Duration) (response *tikvrpc.Response, err error) {
+		select {
+		case <-ctx.Done():
+			return nil, errors.New("timeout")
+		default:
+		}
+		// Return `mismatch peer id` when accesses the leader.
+		if addr == s.cluster.GetStore(s.storeIDs[0]).Address {
+			return &tikvrpc.Response{Resp: &kvrpcpb.GetResponse{RegionError: &errorpb.Error{
+				MismatchPeerId: &errorpb.MismatchPeerId{
+					RequestPeerId: 1,
+					StorePeerId:   2,
+				},
+			}}}, nil
+		}
 		return &tikvrpc.Response{Resp: &kvrpcpb.GetResponse{
 			Value: []byte(addr),
 		}}, nil
@@ -1585,9 +1599,11 @@ func (s *testRegionRequestToThreeStoresSuite) TestDoNotTryUnreachableLeader() {
 	regionStore := region.getStore()
 	leader, _, _, _ := region.WorkStorePeer(regionStore)
 	follower, _, _, _ := region.FollowerStorePeer(regionStore, 0, &storeSelectorOp{})
+	var numCalls int
 
 	s.regionRequestSender.client = &fnClient{fn: func(ctx context.Context, addr string, req *tikvrpc.Request, timeout time.Duration) (response *tikvrpc.Response, err error) {
-		if req.StaleRead && addr == follower.addr {
+		if req.StaleRead && addr == follower.addr && numCalls == 0 {
+			numCalls++
 			return &tikvrpc.Response{Resp: &kvrpcpb.GetResponse{RegionError: &errorpb.Error{DataIsNotReady: &errorpb.DataIsNotReady{}}}}, nil
 		}
 		return &tikvrpc.Response{Resp: &kvrpcpb.GetResponse{