substitute with function

Author: chiacyu

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

@@ -279,10 +279,7 @@ public void testWriteTimout() throws Exception {
 
     setUpMockRequestProto(context);
     AtomicReference<Throwable> throwable = new AtomicReference<>(null);
-    Function<Throwable, ? extends Message> throwableSetter = t -> {
-      throwable.set(t);
-      return null;
-    };
+    Function<Throwable, ? extends Message> throwableSetter = getThrowableSetter(throwable);
     Field writeChunkWaitMaxNs = stateMachine.getClass().getDeclaredField("writeChunkWaitMaxNs");
     writeChunkWaitMaxNs.setAccessible(true);
     writeChunkWaitMaxNs.set(stateMachine, 1000_000_000);