diff --git a/hadoop-ozone/client/src/main/java/org/apache/hadoop/ozone/client/io/KeyInputStream.java b/hadoop-ozone/client/src/main/java/org/apache/hadoop/ozone/client/io/KeyInputStream.java
index c11514f8fd42..33db9a7b0ddf 100644
--- a/hadoop-ozone/client/src/main/java/org/apache/hadoop/ozone/client/io/KeyInputStream.java
+++ b/hadoop-ozone/client/src/main/java/org/apache/hadoop/ozone/client/io/KeyInputStream.java
@@ -36,6 +36,7 @@
 import org.apache.hadoop.ozone.OzoneConsts;
 import org.apache.hadoop.ozone.om.helpers.OmKeyInfo;
 import org.apache.hadoop.ozone.om.helpers.OmKeyLocationInfo;
+import org.apache.hadoop.ozone.om.helpers.OmKeyLocationInfoGroup;
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
 
@@ -104,10 +105,16 @@ private static List<BlockExtendedInputStream> createStreams(
     return partStreams;
   }
 
-  private static BlockLocationInfo getBlockLocationInfo(OmKeyInfo newKeyInfo,
+  static BlockLocationInfo getBlockLocationInfo(OmKeyInfo newKeyInfo,
       BlockID blockID) {
-    for (List<OmKeyLocationInfo> locationList
-        : newKeyInfo.getLatestVersionLocations().getLocationLists()) {
+    if (newKeyInfo == null) {
+      return null;
+    }
+    OmKeyLocationInfoGroup latestLocations = newKeyInfo.getLatestVersionLocations();
+    if (latestLocations == null) {
+      return null;
+    }
+    for (List<OmKeyLocationInfo> locationList : latestLocations.getLocationLists()) {
       for (OmKeyLocationInfo location : locationList) {
         if (location.getBlockID().equals(blockID)) {
           return location;
diff --git a/hadoop-ozone/client/src/test/java/org/apache/hadoop/ozone/client/io/TestKeyInputStream.java b/hadoop-ozone/client/src/test/java/org/apache/hadoop/ozone/client/io/TestKeyInputStream.java
new file mode 100644
index 000000000000..43f16addcb16
--- /dev/null
+++ b/hadoop-ozone/client/src/test/java/org/apache/hadoop/ozone/client/io/TestKeyInputStream.java
@@ -0,0 +1,62 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements. See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License. You may obtain a copy of the License at
+ *
+ *      http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package org.apache.hadoop.ozone.client.io;
+
+import static org.junit.jupiter.api.Assertions.assertNull;
+import static org.junit.jupiter.api.Assertions.assertSame;
+import static org.mockito.Mockito.mock;
+import static org.mockito.Mockito.when;
+
+import java.util.Collections;
+import org.apache.hadoop.hdds.client.BlockID;
+import org.apache.hadoop.ozone.om.helpers.OmKeyInfo;
+import org.apache.hadoop.ozone.om.helpers.OmKeyLocationInfo;
+import org.apache.hadoop.ozone.om.helpers.OmKeyLocationInfoGroup;
+import org.junit.jupiter.api.Test;
+
+/**
+ * Test {@link KeyInputStream#getBlockLocationInfo}.
+ */
+public class TestKeyInputStream {
+
+  private static final BlockID BLOCK_ID = new BlockID(1L, 1L);
+
+  @Test
+  public void testGetBlockLocationInfoWithNullKeyInfo() {
+    assertNull(KeyInputStream.getBlockLocationInfo(null, BLOCK_ID));
+  }
+
+  @Test
+  public void testGetBlockLocationInfoWithoutLocationVersions() {
+    OmKeyInfo keyInfo = mock(OmKeyInfo.class);
+    when(keyInfo.getLatestVersionLocations()).thenReturn(null);
+
+    assertNull(KeyInputStream.getBlockLocationInfo(keyInfo, BLOCK_ID));
+  }
+
+  @Test
+  public void testGetBlockLocationInfoFindsBlock() {
+    OmKeyLocationInfo location = new OmKeyLocationInfo.Builder().setBlockID(BLOCK_ID).build();
+    OmKeyInfo keyInfo = mock(OmKeyInfo.class);
+    when(keyInfo.getLatestVersionLocations())
+        .thenReturn(new OmKeyLocationInfoGroup(0, Collections.singletonList(location)));
+
+    assertSame(location, KeyInputStream.getBlockLocationInfo(keyInfo, BLOCK_ID));
+    assertNull(KeyInputStream.getBlockLocationInfo(keyInfo, new BlockID(2L, 2L)));
+  }
+}