test(lineage): Add test for scroll across lineage (#8728)

Co-authored-by: Indy Prentice <[email protected]>
datahub-project · Aug 28, 2023 · 97019d8 · 97019d8
1 parent 2f11f24
commit 97019d8
Show file tree

Hide file tree

Showing 3 changed files with 91 additions and 1 deletion.
diff --git a/metadata-io/src/main/java/com/linkedin/metadata/search/LineageSearchService.java b/metadata-io/src/main/java/com/linkedin/metadata/search/LineageSearchService.java
@@ -598,7 +598,9 @@ public LineageScrollResult scrollAcrossLineage(@Nonnull Urn sourceUrn, @Nonnull
     List<LineageRelationship> lineageRelationships =
         filterRelationships(lineageResult, new HashSet<>(entities), inputFilters);
 
-    return getScrollResultInBatches(lineageRelationships, input != null ? input : "*", inputFilters, sortCriterion,
+    Filter reducedFilters =
+        SearchUtils.removeCriteria(inputFilters, criterion -> criterion.getField().equals(DEGREE_FILTER_INPUT));
+    return getScrollResultInBatches(lineageRelationships, input != null ? input : "*", reducedFilters, sortCriterion,
         scrollId, keepAlive, size, searchFlags);
   }
 

diff --git a/metadata-io/src/test/java/com/linkedin/metadata/search/LineageSearchServiceTest.java b/metadata-io/src/test/java/com/linkedin/metadata/search/LineageSearchServiceTest.java
@@ -317,6 +317,83 @@ public void testSearchService() throws Exception {
 
   }
 
+  @Test
+  public void testScrollAcrossLineage() throws Exception {
+    when(_graphService.getLineage(eq(TEST_URN), eq(LineageDirection.DOWNSTREAM), anyInt(), anyInt(),
+        anyInt(), eq(null), eq(null))).thenReturn(mockResult(Collections.emptyList()));
+    LineageScrollResult scrollResult = scrollAcrossLineage(null, TEST1);
+    assertEquals(scrollResult.getNumEntities().intValue(), 0);
+    assertNull(scrollResult.getScrollId());
+    scrollResult = scrollAcrossLineage(null, TEST1);
+    assertEquals(scrollResult.getNumEntities().intValue(), 0);
+    assertNull(scrollResult.getScrollId());
+    clearCache(false);
+
+    when(_graphService.getLineage(eq(TEST_URN), eq(LineageDirection.DOWNSTREAM), anyInt(), anyInt(),
+        anyInt(), eq(null), eq(null))).thenReturn(
+        mockResult(ImmutableList.of(new LineageRelationship().setEntity(TEST_URN).setType("test").setDegree(1))));
+    // just testing null input does not throw any exception
+    scrollAcrossLineage(null, null);
+
+    scrollResult = scrollAcrossLineage(null, TEST);
+    assertEquals(scrollResult.getNumEntities().intValue(), 0);
+    assertNull(scrollResult.getScrollId());
+    scrollResult = scrollAcrossLineage(null, TEST1);
+    assertEquals(scrollResult.getNumEntities().intValue(), 0);
+    assertNull(scrollResult.getScrollId());
+    clearCache(false);
+
+    Urn urn = new TestEntityUrn("test1", "urn1", "VALUE_1");
+    ObjectNode document = JsonNodeFactory.instance.objectNode();
+    document.set("urn", JsonNodeFactory.instance.textNode(urn.toString()));
+    document.set("keyPart1", JsonNodeFactory.instance.textNode("test"));
+    document.set("textFieldOverride", JsonNodeFactory.instance.textNode("textFieldOverride"));
+    document.set("browsePaths", JsonNodeFactory.instance.textNode("/a/b/c"));
+    _elasticSearchService.upsertDocument(ENTITY_NAME, document.toString(), urn.toString());
+    syncAfterWrite(_bulkProcessor);
+
+    when(_graphService.getLineage(eq(TEST_URN), eq(LineageDirection.DOWNSTREAM), anyInt(), anyInt(),
+        anyInt(), eq(null), eq(null))).thenReturn(mockResult(Collections.emptyList()));
+    scrollResult = scrollAcrossLineage(null, TEST1);
+    assertEquals(scrollResult.getNumEntities().intValue(), 0);
+    assertEquals(scrollResult.getEntities().size(), 0);
+    assertNull(scrollResult.getScrollId());
+    clearCache(false);
+
+    when(_graphService.getLineage(eq(TEST_URN), eq(LineageDirection.DOWNSTREAM), anyInt(), anyInt(),
+        anyInt(), eq(null), eq(null))).thenReturn(
+        mockResult(ImmutableList.of(new LineageRelationship().setEntity(urn).setType("test").setDegree(1))));
+    scrollResult = scrollAcrossLineage(null, TEST1);
+    assertEquals(scrollResult.getNumEntities().intValue(), 1);
+    assertEquals(scrollResult.getEntities().get(0).getEntity(), urn);
+    assertEquals(scrollResult.getEntities().get(0).getDegree().intValue(), 1);
+    assertNull(scrollResult.getScrollId());
+
+    scrollResult = scrollAcrossLineage(QueryUtils.newFilter("degree.keyword", "1"), TEST1);
+    assertEquals(scrollResult.getNumEntities().intValue(), 1);
+    assertEquals(scrollResult.getEntities().get(0).getEntity(), urn);
+    assertEquals(scrollResult.getEntities().get(0).getDegree().intValue(), 1);
+    assertNull(scrollResult.getScrollId());
+
+    scrollResult = scrollAcrossLineage(QueryUtils.newFilter("degree.keyword", "2"), TEST1);
+    assertEquals(scrollResult.getNumEntities().intValue(), 0);
+    assertEquals(scrollResult.getEntities().size(), 0);
+    assertNull(scrollResult.getScrollId());
+    clearCache(false);
+
+    // Cleanup
+    _elasticSearchService.deleteDocument(ENTITY_NAME, urn.toString());
+    syncAfterWrite(_bulkProcessor);
+
+    when(_graphService.getLineage(eq(TEST_URN), eq(LineageDirection.DOWNSTREAM), anyInt(), anyInt(),
+        anyInt())).thenReturn(
+        mockResult(ImmutableList.of(new LineageRelationship().setEntity(urn).setType("test1").setDegree(1))));
+    scrollResult = scrollAcrossLineage(null, TEST1);
+
+    assertEquals(scrollResult.getNumEntities().intValue(), 0);
+    assertNull(scrollResult.getScrollId());
+  }
+
   @Test
   public void testLightningSearchService() throws Exception {
     // Mostly this test ensures the code path is exercised
@@ -731,6 +808,16 @@ private LineageSearchResult searchAcrossLineage(@Nullable Filter filter, @Nullab
         new SearchFlags().setSkipCache(true));
   }
 
+  private LineageScrollResult scrollAcrossLineage(@Nullable Filter filter, @Nullable String input, String scrollId, int size) {
+    return _lineageSearchService.scrollAcrossLineage(TEST_URN, LineageDirection.DOWNSTREAM, ImmutableList.of(), input,
+        null, filter, null, scrollId,  "5m",  size, null, null,
+        new SearchFlags().setSkipCache(true));
+  }
+
+  private LineageScrollResult scrollAcrossLineage(@Nullable Filter filter, @Nullable String input) {
+    return scrollAcrossLineage(filter, input, null, 10);
+  }
+
   @Test
   public void testCanDoLightning() throws Exception {
     Map<String, Integer> platformCounts = new HashMap<>();

diff --git a/...test/java/com/linkedin/metadata/search/elasticsearch/fixtures/SampleDataFixtureTests.java b/...test/java/com/linkedin/metadata/search/elasticsearch/fixtures/SampleDataFixtureTests.java
@@ -161,6 +161,7 @@ public void testSearchFieldConfig() throws IOException {
                     // this is a subfield therefore cannot have a subfield
                     assertFalse(test.hasKeywordSubfield());
                     assertFalse(test.hasDelimitedSubfield());
+                    assertFalse(test.hasWordGramSubfields());
 
                     String[] fieldAndSubfield = test.fieldName().split("[.]", 2);