diff --git a/hadoop-ozone/client/src/test/java/org/apache/hadoop/ozone/client/MockOmTransport.java b/hadoop-ozone/client/src/test/java/org/apache/hadoop/ozone/client/MockOmTransport.java
index 36ab47af7e44..17a7f6b28ca3 100644
--- a/hadoop-ozone/client/src/test/java/org/apache/hadoop/ozone/client/MockOmTransport.java
+++ b/hadoop-ozone/client/src/test/java/org/apache/hadoop/ozone/client/MockOmTransport.java
@@ -19,6 +19,7 @@
 
 import org.apache.hadoop.io.Text;
 import org.apache.hadoop.ozone.om.protocolPB.OmTransport;
+import org.apache.hadoop.ozone.protocol.proto.OzoneManagerProtocolProtos;
 import org.apache.hadoop.ozone.protocol.proto.OzoneManagerProtocolProtos.BucketInfo;
 import org.apache.hadoop.ozone.protocol.proto.OzoneManagerProtocolProtos.CommitKeyRequest;
 import org.apache.hadoop.ozone.protocol.proto.OzoneManagerProtocolProtos.CommitKeyResponse;
@@ -116,6 +117,9 @@ public OMResponse submitRequest(OMRequest payload) throws IOException {
       return response(payload,
           r -> r.setServiceListResponse(
               serviceList(payload.getServiceListRequest())));
+    case AllocateBlock:
+      return response(payload, r -> r.setAllocateBlockResponse(
+          allocateBlock(payload.getAllocateBlockRequest())));
     default:
       throw new IllegalArgumentException(
           "Mock version of om call " + payload.getCmdType()
@@ -123,6 +127,14 @@ public OMResponse submitRequest(OMRequest payload) throws IOException {
     }
   }
 
+  private OzoneManagerProtocolProtos.AllocateBlockResponse allocateBlock(
+      OzoneManagerProtocolProtos.AllocateBlockRequest allocateBlockRequest) {
+    return OzoneManagerProtocolProtos.AllocateBlockResponse.newBuilder()
+        .setKeyLocation(
+            blockAllocator.allocateBlock(allocateBlockRequest.getKeyArgs())
+                .iterator().next()).build();
+  }
+
   private DeleteVolumeResponse deleteVolume(
       DeleteVolumeRequest deleteVolumeRequest) {
     volumes.remove(deleteVolumeRequest.getVolumeName());
diff --git a/hadoop-ozone/client/src/test/java/org/apache/hadoop/ozone/client/TestOzoneClient.java b/hadoop-ozone/client/src/test/java/org/apache/hadoop/ozone/client/TestOzoneClient.java
index ccc58982a6d8..2f4ebcfb925e 100644
--- a/hadoop-ozone/client/src/test/java/org/apache/hadoop/ozone/client/TestOzoneClient.java
+++ b/hadoop-ozone/client/src/test/java/org/apache/hadoop/ozone/client/TestOzoneClient.java
@@ -142,15 +142,9 @@ public void testCreateBucket()
 
   @Test
   public void testPutKeyRatisOneNode() throws IOException {
-    String volumeName = UUID.randomUUID().toString();
-    String bucketName = UUID.randomUUID().toString();
     Instant testStartTime = Instant.now();
-
     String value = "sample value";
-    store.createVolume(volumeName);
-    OzoneVolume volume = store.getVolume(volumeName);
-    volume.createBucket(bucketName);
-    OzoneBucket bucket = volume.getBucket(bucketName);
+    OzoneBucket bucket = getOzoneBucket();
 
     for (int i = 0; i < 10; i++) {
       String keyName = UUID.randomUUID().toString();
@@ -171,4 +165,30 @@ public void testPutKeyRatisOneNode() throws IOException {
       Assert.assertFalse(key.getModificationTime().isBefore(testStartTime));
     }
   }
+
+  @Test
+  public void testPutKeyAllocateBlock() throws IOException {
+    String value = new String(new byte[1024], UTF_8);
+    OzoneBucket bucket = getOzoneBucket();
+
+    for (int i = 0; i < 10; i++) {
+      String keyName = UUID.randomUUID().toString();
+
+      try (OzoneOutputStream out = bucket
+          .createKey(keyName, value.getBytes(UTF_8).length,
+              ReplicationType.RATIS, ONE, new HashMap<>())) {
+        out.write(value.getBytes(UTF_8));
+        out.write(value.getBytes(UTF_8));
+      }
+    }
+  }
+
+  private OzoneBucket getOzoneBucket() throws IOException {
+    String volumeName = UUID.randomUUID().toString();
+    String bucketName = UUID.randomUUID().toString();
+    store.createVolume(volumeName);
+    OzoneVolume volume = store.getVolume(volumeName);
+    volume.createBucket(bucketName);
+    return volume.getBucket(bucketName);
+  }
 }
\ No newline at end of file