fix unit test

Signed-off-by: iosmanthus <myosmanthustree@gmail.com>

src/test/java/org/tikv/common/MockThreeStoresTest.java

@@ -30,6 +30,7 @@ import org.tikv.common.region.TiRegion;
 import org.tikv.common.region.TiStore;
 import org.tikv.kvproto.Metapb;
 import org.tikv.kvproto.Pdpb;
+import org.tikv.kvproto.Pdpb.GetAllStoresResponse;
 
 public class MockThreeStoresTest extends PDMockServerTest {
 
@@ -68,17 +69,17 @@ public class MockThreeStoresTest extends PDMockServerTest {
         ImmutableList.of(
             Metapb.Store.newBuilder()
                 .setAddress("127.0.0.1:" + ports[0])
-                .setVersion("5.0.0")
+                .setVersion(Version.API_V2)
                 .setId(0x1)
                 .build(),
             Metapb.Store.newBuilder()
                 .setAddress("127.0.0.1:" + ports[1])
-                .setVersion("5.0.0")
+                .setVersion(Version.API_V2)
                 .setId(0x2)
                 .build(),
             Metapb.Store.newBuilder()
                 .setAddress("127.0.0.1:" + ports[2])
-                .setVersion("5.0.0")
+                .setVersion(Version.API_V2)
                 .setId(0x3)
                 .build());
 
@@ -94,6 +95,10 @@ public class MockThreeStoresTest extends PDMockServerTest {
             int i = (int) request.getStoreId() - 1;
             return Pdpb.GetStoreResponse.newBuilder().setStore(stores.get(i)).build();
           });
+      server.addGetAllStoresListener(
+          request -> {
+            return GetAllStoresResponse.newBuilder().addAllStores(stores).build();
+          });
     }
 
     this.region =

src/test/java/org/tikv/common/PDClientV2MockTest.java

@@ -17,6 +17,7 @@
 
 package org.tikv.common;
 
+import com.google.common.collect.ImmutableList;
 import com.google.protobuf.ByteString;
 import java.util.List;
 import org.junit.Assert;
@@ -27,7 +28,10 @@ import org.tikv.common.codec.CodecDataOutput;
 import org.tikv.common.util.ConcreteBackOffer;
 import org.tikv.common.util.Pair;
 import org.tikv.kvproto.Metapb;
+import org.tikv.kvproto.Metapb.Store;
+import org.tikv.kvproto.Metapb.StoreState;
 import org.tikv.kvproto.Pdpb;
+import org.tikv.kvproto.Pdpb.GetAllStoresResponse;
 import org.tikv.kvproto.Pdpb.GetRegionResponse;
 import org.tikv.kvproto.Pdpb.Region;
 import org.tikv.kvproto.Pdpb.ScanRegionsResponse;
@@ -35,6 +39,18 @@ import org.tikv.kvproto.Pdpb.ScanRegionsResponse;
 public class PDClientV2MockTest extends PDMockServerTest {
   @Before
   public void init() throws Exception {
+    leader.addGetAllStoresListener(
+        request -> {
+          return GetAllStoresResponse.newBuilder()
+              .addAllStores(
+                  ImmutableList.of(
+                      Store.newBuilder()
+                          .setId(0x1)
+                          .setState(StoreState.Up)
+                          .setVersion(Version.API_V2)
+                          .build()))
+              .build();
+        });
     upgradeToV2Cluster();
   }
 

src/test/java/org/tikv/common/PDMockServer.java

@@ -30,6 +30,8 @@ import java.net.ServerSocket;
 import java.util.Optional;
 import java.util.function.Function;
 import org.tikv.kvproto.PDGrpc;
+import org.tikv.kvproto.Pdpb.GetAllStoresRequest;
+import org.tikv.kvproto.Pdpb.GetAllStoresResponse;
 import org.tikv.kvproto.Pdpb.GetMembersRequest;
 import org.tikv.kvproto.Pdpb.GetMembersResponse;
 import org.tikv.kvproto.Pdpb.GetRegionByIDRequest;
@@ -54,6 +56,8 @@ public class PDMockServer extends PDGrpc.PDImplBase {
 
   private Function<ScanRegionsRequest, ScanRegionsResponse> scanRegionsListener;
 
+  private Function<GetAllStoresRequest, GetAllStoresResponse> getAllStoresListener;
+
   public void addGetMembersListener(Function<GetMembersRequest, GetMembersResponse> func) {
     getMembersListener = func;
   }
@@ -144,6 +148,20 @@ public class PDMockServer extends PDGrpc.PDImplBase {
     }
   }
 
+  public void addGetAllStoresListener(Function<GetAllStoresRequest, GetAllStoresResponse> func) {
+    getAllStoresListener = func;
+  }
+
+  @Override
+  public void getAllStores(GetAllStoresRequest request, StreamObserver<GetAllStoresResponse> resp) {
+    try {
+      resp.onNext(Optional.ofNullable(getAllStoresListener.apply(request)).get());
+      resp.onCompleted();
+    } catch (Exception e) {
+      resp.onError(Status.INTERNAL.asRuntimeException());
+    }
+  }
+
   public void start(long clusterId) throws IOException {
     int port;
     try (ServerSocket s = new ServerSocket(0)) {