diff --git a/src/main/java/org/opensearch/neuralsearch/util/ProcessorDocumentUtils.java b/src/main/java/org/opensearch/neuralsearch/util/ProcessorDocumentUtils.java index 6f9297e5c..0cbf4534d 100644 --- a/src/main/java/org/opensearch/neuralsearch/util/ProcessorDocumentUtils.java +++ b/src/main/java/org/opensearch/neuralsearch/util/ProcessorDocumentUtils.java @@ -296,30 +296,6 @@ private static void unflattenSingleItem(String key, Object value, Map searchResponseAsMap = search(INDEX_NAME, queryNestedHighLevel, 2); assertNotNull(searchResponseAsMap); - Map hits = (Map) searchResponseAsMap.get("hits"); - assertNotNull(hits); - - List> listOfHits = (List>) hits.get("hits"); - assertNotNull(listOfHits); - assertEquals(1, listOfHits.size()); + assertEquals(1, getHitCount(searchResponseAsMap)); - Map innerHitDetails = listOfHits.getFirst(); + Map innerHitDetails = getFirstInnerHit(searchResponseAsMap); assertEquals("5", innerHitDetails.get("_id")); } finally { wipeOfTestResources(INDEX_NAME, PIPELINE_NAME, modelId, null);