refactor functions

Author: chiacyu

hadoop-hdds/container-service/src/test/java/org/apache/hadoop/ozone/container/common/transport/server/ratis/TestContainerStateMachine.java

@@ -131,7 +131,7 @@ public void testWriteFailure(boolean failWithException) throws ExecutionExceptio
     when(trx.getStateMachineContext()).thenReturn(context);
 
     setUpMockDispatcherReturn(failWithException);
-    setUpMockRequestProto(context);
+    setUpMockRequestProtoReturn(context, "Test Data", 1, 1);
 
     AtomicReference<Throwable> throwable = new AtomicReference<>(null);
     Function<Throwable, ? extends Message> throwableSetter = getThrowableSetter(throwable);
@@ -144,13 +144,7 @@ public void testWriteFailure(boolean failWithException) throws ExecutionExceptio
     assertResults(failWithException, throwable);
 
     // Writing data to another container(containerId 2) should also fail.
-    when(context.getRequestProto()).thenReturn(ContainerProtos.ContainerCommandRequestProto.newBuilder()
-        .setCmdType(ContainerProtos.Type.WriteChunk).setWriteChunk(
-            ContainerProtos.WriteChunkRequestProto.newBuilder().setData(ByteString.copyFromUtf8("Test Data"))
-                .setBlockID(
-                    ContainerProtos.DatanodeBlockID.newBuilder().setContainerID(2).setLocalID(1).build()).build())
-        .setContainerID(2)
-        .setDatanodeUuid(UUID.randomUUID().toString()).build());
+    setUpMockRequestProtoReturn(context, "Test Data", 2, 1);
     stateMachine.write(entryNext, trx).exceptionally(throwableSetter).get();
     verify(dispatcher, times(0)).dispatch(any(ContainerProtos.ContainerCommandRequestProto.class),
         any(DispatcherContext.class));
@@ -170,13 +164,15 @@ public final void setUpMockDispatcherReturn(boolean failWithException) {
     }
   }
 
-  public final void setUpMockRequestProto(ContainerStateMachine.Context context) {
+  public final void setUpMockRequestProtoReturn(ContainerStateMachine.Context context, String content,
+                                                int containerId, int localId) {
     when(context.getRequestProto()).thenReturn(ContainerProtos.ContainerCommandRequestProto.newBuilder()
         .setCmdType(ContainerProtos.Type.WriteChunk).setWriteChunk(
-            ContainerProtos.WriteChunkRequestProto.newBuilder().setData(ByteString.copyFromUtf8("Test Data"))
+            ContainerProtos.WriteChunkRequestProto.newBuilder().setData(ByteString.copyFromUtf8(content))
                 .setBlockID(
-                    ContainerProtos.DatanodeBlockID.newBuilder().setContainerID(1).setLocalID(1).build()).build())
-        .setContainerID(1)
+                    ContainerProtos.DatanodeBlockID.newBuilder().setContainerID(containerId)
+                        .setLocalID(localId).build()).build())
+        .setContainerID(containerId)
         .setDatanodeUuid(UUID.randomUUID().toString()).build());
   }
 
@@ -198,6 +194,16 @@ public final void assertResults(boolean failWithException, AtomicReference<Throw
     }
   }
 
+  public final void setUpLogProtoReturn(ContainerStateMachine.Context context, int containerId, int localId) {
+    when(context.getLogProto()).thenReturn(ContainerProtos.ContainerCommandRequestProto.newBuilder()
+        .setCmdType(ContainerProtos.Type.WriteChunk).setWriteChunk(
+            ContainerProtos.WriteChunkRequestProto.newBuilder().setBlockID(
+                ContainerProtos.DatanodeBlockID.newBuilder().
+                    setContainerID(containerId).setLocalID(localId).build()).build())
+        .setContainerID(containerId)
+        .setDatanodeUuid(UUID.randomUUID().toString()).build());
+  }
+
   @ParameterizedTest
   @ValueSource(booleans = {true, false})
   public void testApplyTransactionFailure(boolean failWithException) throws ExecutionException,
@@ -212,12 +218,7 @@ public void testApplyTransactionFailure(boolean failWithException) throws Execut
 
     setUpMockDispatcherReturn(failWithException);
     // Failing apply transaction on congtainer 1.
-    when(context.getLogProto()).thenReturn(ContainerProtos.ContainerCommandRequestProto.newBuilder()
-        .setCmdType(ContainerProtos.Type.WriteChunk).setWriteChunk(
-            ContainerProtos.WriteChunkRequestProto.newBuilder().setBlockID(
-                ContainerProtos.DatanodeBlockID.newBuilder().setContainerID(1).setLocalID(1).build()).build())
-        .setContainerID(1)
-        .setDatanodeUuid(UUID.randomUUID().toString()).build());
+    setUpLogProtoReturn(context, 1, 1);
     AtomicReference<Throwable> throwable = new AtomicReference<>(null);
     Function<Throwable, ? extends Message> throwableSetter = getThrowableSetter(throwable);
     //apply transaction will fail because of runtime exception thrown by dispatcher, which marks the first
@@ -238,13 +239,7 @@ public void testApplyTransactionFailure(boolean failWithException) throws Execut
 
     // Another apply transaction on a different container 2 shouldn't fail because the previous apply transaction
     // failure was only on container 1.
-    when(context.getLogProto()).thenReturn(ContainerProtos.ContainerCommandRequestProto.newBuilder()
-        .setCmdType(ContainerProtos.Type.WriteChunk).setWriteChunk(
-            ContainerProtos.WriteChunkRequestProto.newBuilder().setBlockID(
-                ContainerProtos.DatanodeBlockID.newBuilder().setContainerID(2).setLocalID(1).build()).build())
-        .setContainerID(2)
-        .setDatanodeUuid(UUID.randomUUID().toString()).build());
-
+    setUpLogProtoReturn(context, 2, 1);
     reset(dispatcher);
     throwable.set(null);
     when(dispatcher.dispatch(any(), any())).thenReturn(ContainerProtos.ContainerCommandResponseProto
@@ -280,7 +275,7 @@ public void testWriteTimout() throws Exception {
       return null;
     }).when(dispatcher).dispatch(any(), any());
 
-    setUpMockRequestProto(context);
+    setUpMockRequestProtoReturn(context, "Test data", 1, 1);
     AtomicReference<Throwable> throwable = new AtomicReference<>(null);
     Function<Throwable, ? extends Message> throwableSetter = t -> {
       throwable.set(t);