Business Attribute : SearchableRef Unit Test

li-utils/src/main/java/com/linkedin/metadata/Constants.java

@@ -378,8 +378,7 @@ public class Constants {
   public static final String BUSINESS_ATTRIBUTE_KEY_ASPECT_NAME = "businessAttributeKey";
   public static final String BUSINESS_ATTRIBUTE_INFO_ASPECT_NAME = "businessAttributeInfo";
   public static final String BUSINESS_ATTRIBUTE_ASSOCIATION = "businessAttributeAssociation";
-  public static final List<String> SKIP_REFERENCE_ASPECT =
-      Arrays.asList("ownership", "status", "institutionalMemory");
+  public static final List<String> SKIP_REFERENCE_ASPECT = List.of("ownership", "status", "institutionalMemory");
 
   // Posts
   public static final String POST_INFO_ASPECT_NAME = "postInfo";

metadata-io/src/test/java/com/linkedin/metadata/search/elasticsearch/query/request/TestSearchFieldConfig.java

@@ -0,0 +1,64 @@
+package com.linkedin.metadata.search.elasticsearch.query.request;
+
+import com.linkedin.metadata.models.SearchableRefFieldSpec;
+import com.linkedin.metadata.models.registry.ConfigEntityRegistry;
+import com.linkedin.metadata.models.registry.EntityRegistry;
+import java.util.Optional;
+import java.util.Set;
+import org.junit.jupiter.api.Assertions;
+import org.testng.annotations.Test;
+
+@Test
+public class TestSearchFieldConfig {
+
+  void setup() {}
+
+  /**
+   *
+   *
+   * <ul>
+   *   <li>{@link SearchFieldConfig#detectSubFieldType( SearchableRefFieldSpec, int, EntityRegistry
+   *       ) }
+   * </ul>
+   */
+  @Test
+  public void detectSubFieldType() {
+    EntityRegistry entityRegistry = getTestEntityRegistry();
+    SearchableRefFieldSpec searchableRefFieldSpec =
+        entityRegistry.getEntitySpec("testRefEntity").getSearchableRefFieldSpecs().get(0);
+
+    Set<SearchFieldConfig> responseForNonZeroDepth =
+        SearchFieldConfig.detectSubFieldType(searchableRefFieldSpec, 1, entityRegistry);
+    Assertions.assertTrue(
+        responseForNonZeroDepth.stream()
+            .anyMatch(
+                searchFieldConfig ->
+                    searchFieldConfig.fieldName().equals("refEntityUrns.displayName")));
+    Assertions.assertTrue(
+        responseForNonZeroDepth.stream()
+            .anyMatch(
+                searchFieldConfig -> searchFieldConfig.fieldName().equals("refEntityUrns.urn")));
+    Assertions.assertTrue(
+        responseForNonZeroDepth.stream()
+            .anyMatch(
+                searchFieldConfig ->
+                    searchFieldConfig.fieldName().equals("refEntityUrns.editedFieldDescriptions")));
+
+    Set<SearchFieldConfig> responseForZeroDepth =
+        SearchFieldConfig.detectSubFieldType(searchableRefFieldSpec, 0, entityRegistry);
+    Optional<SearchFieldConfig> searchFieldConfigToCompare =
+        responseForZeroDepth.stream()
+            .filter(searchFieldConfig -> searchFieldConfig.fieldName().equals("refEntityUrns"))
+            .findFirst();
+
+    Assertions.assertTrue(searchFieldConfigToCompare.isPresent());
+    Assertions.assertEquals("query_urn_component", searchFieldConfigToCompare.get().analyzer());
+  }
+
+  private EntityRegistry getTestEntityRegistry() {
+    return new ConfigEntityRegistry(
+        TestSearchFieldConfig.class
+            .getClassLoader()
+            .getResourceAsStream("test-entity-registry.yaml"));
+  }
+}

metadata-io/src/test/java/com/linkedin/metadata/search/indexbuilder/MappingsBuilderTest.java

@@ -3,12 +3,19 @@
 import static com.linkedin.metadata.Constants.*;
 import static org.testng.Assert.*;
 
+import com.datahub.test.TestRefEntity;
 import com.google.common.collect.ImmutableMap;
 import com.linkedin.common.UrnArray;
 import com.linkedin.common.urn.Urn;
 import com.linkedin.metadata.TestEntitySpecBuilder;
+import com.linkedin.metadata.models.EntitySpec;
+import com.linkedin.metadata.models.EntitySpecBuilder;
+import com.linkedin.metadata.models.registry.ConfigEntityRegistry;
+import com.linkedin.metadata.models.registry.EntityRegistry;
 import com.linkedin.metadata.search.elasticsearch.indexbuilder.MappingsBuilder;
+import com.linkedin.metadata.search.elasticsearch.query.request.TestSearchFieldConfig;
 import com.linkedin.structured.StructuredPropertyDefinition;
+import java.io.Serializable;
 import java.net.URISyntaxException;
 import java.util.List;
 import java.util.Map;
@@ -271,4 +278,61 @@ public void testGetMappingsForStructuredProperty() throws URISyntaxException {
     mappings = structuredPropertyFieldMappingsNumber.get(keyInMap);
     assertEquals(Map.of("type", "double"), mappings);
   }
+
+  @Test
+  public void testRefMappingsBuilder() {
+    EntityRegistry entityRegistry = getTestEntityRegistry();
+    MappingsBuilder.setEntityRegistry(entityRegistry);
+    EntitySpec entitySpec = new EntitySpecBuilder().buildEntitySpec(new TestRefEntity().schema());
+    Map<String, Object> result = MappingsBuilder.getMappings(entitySpec);
+    assertEquals(result.size(), 1);
+    Map<String, Object> properties = (Map<String, Object>) result.get("properties");
+    assertEquals(properties.size(), 6);
+    ImmutableMap<String, Serializable> expectedURNField =
+        ImmutableMap.of(
+            "type",
+            "keyword",
+            "fields",
+            ImmutableMap.of(
+                "delimited",
+                ImmutableMap.of(
+                    "type",
+                    "text",
+                    "analyzer",
+                    "urn_component",
+                    "search_analyzer",
+                    "query_urn_component",
+                    "search_quote_analyzer",
+                    "quote_analyzer"),
+                "ngram",
+                ImmutableMap.of(
+                    "type",
+                    "search_as_you_type",
+                    "max_shingle_size",
+                    "4",
+                    "doc_values",
+                    "false",
+                    "analyzer",
+                    "partial_urn_component")));
+    assertEquals(properties.get("urn"), expectedURNField);
+    assertEquals(properties.get("runId"), ImmutableMap.of("type", "keyword"));
+    assertTrue(properties.containsKey("editedFieldDescriptions"));
+    assertTrue(properties.containsKey("displayName"));
+    assertTrue(properties.containsKey("refEntityUrns"));
+    // @SearchableRef Field
+    Map<String, Object> refField = (Map<String, Object>) properties.get("refEntityUrns");
+    assertEquals(refField.size(), 1);
+    Map<String, Object> refFieldProperty = (Map<String, Object>) refField.get("properties");
+
+    assertEquals(refFieldProperty.get("urn"), expectedURNField);
+    assertTrue(refFieldProperty.containsKey("displayName"));
+    assertTrue(refFieldProperty.containsKey("editedFieldDescriptions"));
+  }
+
+  private EntityRegistry getTestEntityRegistry() {
+    return new ConfigEntityRegistry(
+        TestSearchFieldConfig.class
+            .getClassLoader()
+            .getResourceAsStream("test-entity-registry.yaml"));
+  }
 }

metadata-io/src/test/java/com/linkedin/metadata/search/transformer/SearchDocumentTransformerTest.java

@@ -1,6 +1,7 @@
 package com.linkedin.metadata.search.transformer;
 
 import static com.linkedin.metadata.Constants.*;
+import static org.mockito.Mockito.*;
 import static org.testng.Assert.assertEquals;
 import static org.testng.Assert.assertFalse;
 import static org.testng.Assert.assertTrue;
@@ -13,11 +14,22 @@
 import com.fasterxml.jackson.databind.node.JsonNodeFactory;
 import com.fasterxml.jackson.databind.node.JsonNodeType;
 import com.fasterxml.jackson.databind.node.ObjectNode;
+import com.linkedin.common.urn.Urn;
+import com.linkedin.data.DataMapBuilder;
+import com.linkedin.entity.Aspect;
 import com.linkedin.metadata.TestEntitySpecBuilder;
 import com.linkedin.metadata.TestEntityUtil;
+import com.linkedin.metadata.aspect.AspectRetriever;
 import com.linkedin.metadata.models.EntitySpec;
+import com.linkedin.metadata.models.SearchableRefFieldSpec;
+import com.linkedin.metadata.models.registry.ConfigEntityRegistry;
+import com.linkedin.metadata.models.registry.EntityRegistry;
+import com.linkedin.metadata.search.elasticsearch.query.request.TestSearchFieldConfig;
+import com.linkedin.r2.RemoteInvocationException;
 import java.io.IOException;
-import java.util.Optional;
+import java.net.URISyntaxException;
+import java.util.*;
+import org.mockito.Mockito;
 import org.testng.annotations.Test;
 
 public class SearchDocumentTransformerTest {
@@ -132,4 +144,153 @@ public void testTransformMaxFieldValue() throws IOException {
             .add("123")
             .add("0123456789"));
   }
+
+  /**
+   *
+   *
+   * <ul>
+   *   <li>{@link SearchDocumentTransformer#setSearchableRefValue(SearchableRefFieldSpec, List,
+   *       ObjectNode, Boolean ) }
+   * </ul>
+   */
+  @Test
+  public void testSetSearchableRefValue() throws URISyntaxException, RemoteInvocationException {
+    AspectRetriever aspectRetriever = Mockito.mock(AspectRetriever.class);
+    SearchDocumentTransformer searchDocumentTransformer =
+        new SearchDocumentTransformer(1000, 1000, 1000);
+    searchDocumentTransformer.setAspectRetriever(aspectRetriever);
+    EntityRegistry entityRegistry = getTestEntityRegistry();
+    List<Object> urnList = List.of(Urn.createFromString("urn:li:refEntity:1"));
+
+    DataMapBuilder dataMapBuilder = new DataMapBuilder();
+    dataMapBuilder.addKVPair("fieldPath", "refEntityUrn");
+    dataMapBuilder.addKVPair("name", "refEntityUrnName");
+    dataMapBuilder.addKVPair("description", "refEntityUrn1 description details");
+    Aspect aspect = new Aspect(dataMapBuilder.convertToDataMap());
+
+    ObjectNode searchDocument = JsonNodeFactory.instance.objectNode();
+    SearchableRefFieldSpec searchableRefFieldSpec =
+        entityRegistry.getEntitySpec("testRefEntity").getSearchableRefFieldSpecs().get(0);
+
+    // Mock Behaviour
+    Mockito.when(aspectRetriever.getEntityRegistry()).thenReturn(entityRegistry);
+    Mockito.when(aspectRetriever.getLatestAspectObject(any(), anyString())).thenReturn(aspect);
+
+    searchDocumentTransformer.setSearchableRefValue(
+        searchableRefFieldSpec, urnList, searchDocument, false);
+    assertTrue(searchDocument.has("refEntityUrns"));
+    assertEquals(searchDocument.get("refEntityUrns").size(), 3);
+    assertTrue(searchDocument.get("refEntityUrns").has("urn"));
+    assertTrue(searchDocument.get("refEntityUrns").has("editedFieldDescriptions"));
+    assertTrue(searchDocument.get("refEntityUrns").has("displayName"));
+    assertEquals(searchDocument.get("refEntityUrns").get("urn").asText(), "urn:li:refEntity:1");
+    assertEquals(
+        searchDocument.get("refEntityUrns").get("editedFieldDescriptions").asText(),
+        "refEntityUrn1 description details");
+    assertEquals(
+        searchDocument.get("refEntityUrns").get("displayName").asText(), "refEntityUrnName");
+  }
+
+  @Test
+  public void testSetSearchableRefValue_WithNonURNField() throws URISyntaxException {
+    AspectRetriever aspectRetriever = Mockito.mock(AspectRetriever.class);
+    SearchDocumentTransformer searchDocumentTransformer =
+        new SearchDocumentTransformer(1000, 1000, 1000);
+    searchDocumentTransformer.setAspectRetriever(aspectRetriever);
+    EntityRegistry entityRegistry = getTestEntityRegistry();
+    List<Object> urnList = List.of(Urn.createFromString("urn:li:refEntity:1"));
+
+    ObjectNode searchDocument = JsonNodeFactory.instance.objectNode();
+    SearchableRefFieldSpec searchableRefFieldSpecText =
+        entityRegistry.getEntitySpec("testRefEntity").getSearchableRefFieldSpecs().get(1);
+    searchDocumentTransformer.setSearchableRefValue(
+        searchableRefFieldSpecText, urnList, searchDocument, false);
+    assertTrue(searchDocument.isEmpty());
+  }
+
+  @Test
+  public void testSetSearchableRefValue_RemoteInvocationException()
+      throws URISyntaxException, RemoteInvocationException {
+    AspectRetriever aspectRetriever = Mockito.mock(AspectRetriever.class);
+    SearchDocumentTransformer searchDocumentTransformer =
+        new SearchDocumentTransformer(1000, 1000, 1000);
+    searchDocumentTransformer.setAspectRetriever(aspectRetriever);
+    EntityRegistry entityRegistry = getTestEntityRegistry();
+    List<Object> urnList = List.of(Urn.createFromString("urn:li:refEntity:1"));
+
+    Mockito.when(aspectRetriever.getEntityRegistry()).thenReturn(entityRegistry);
+    Mockito.when(
+            aspectRetriever.getLatestAspectObject(
+                eq(Urn.createFromString("urn:li:refEntity:1")), anyString()))
+        .thenThrow(new RemoteInvocationException("Error"));
+
+    ObjectNode searchDocument = JsonNodeFactory.instance.objectNode();
+    SearchableRefFieldSpec searchableRefFieldSpec =
+        entityRegistry.getEntitySpec("testRefEntity").getSearchableRefFieldSpecs().get(0);
+    searchDocumentTransformer.setSearchableRefValue(
+        searchableRefFieldSpec, urnList, searchDocument, false);
+    assertTrue(searchDocument.isEmpty());
+  }
+
+  @Test
+  public void testSetSearchableRefValue_RemoteInvocationException_URNExist()
+      throws URISyntaxException, RemoteInvocationException {
+    AspectRetriever aspectRetriever = Mockito.mock(AspectRetriever.class);
+    SearchDocumentTransformer searchDocumentTransformer =
+        new SearchDocumentTransformer(1000, 1000, 1000);
+    searchDocumentTransformer.setAspectRetriever(aspectRetriever);
+    EntityRegistry entityRegistry = getTestEntityRegistry();
+    List<Object> urnList = List.of(Urn.createFromString("urn:li:refEntity:1"));
+    DataMapBuilder dataMapBuilder = new DataMapBuilder();
+    dataMapBuilder.addKVPair("fieldPath", "refEntityUrn");
+    dataMapBuilder.addKVPair("name", "refEntityUrnName");
+    dataMapBuilder.addKVPair("description", "refEntityUrn1 description details");
+
+    Aspect aspect = new Aspect(dataMapBuilder.convertToDataMap());
+    Mockito.when(aspectRetriever.getEntityRegistry()).thenReturn(entityRegistry);
+    Mockito.when(
+            aspectRetriever.getLatestAspectObject(
+                eq(Urn.createFromString("urn:li:refEntity:1")), anyString()))
+        .thenReturn(aspect)
+        .thenThrow(new RemoteInvocationException("Error"));
+
+    ObjectNode searchDocument = JsonNodeFactory.instance.objectNode();
+    SearchableRefFieldSpec searchableRefFieldSpec =
+        entityRegistry.getEntitySpec("testRefEntity").getSearchableRefFieldSpecs().get(0);
+    searchDocumentTransformer.setSearchableRefValue(
+        searchableRefFieldSpec, urnList, searchDocument, false);
+    assertTrue(searchDocument.has("refEntityUrns"));
+    assertEquals(searchDocument.get("refEntityUrns").size(), 1);
+    assertTrue(searchDocument.get("refEntityUrns").has("urn"));
+    assertEquals(searchDocument.get("refEntityUrns").get("urn").asText(), "urn:li:refEntity:1");
+  }
+
+  @Test
+  void testSetSearchableRefValue_WithInvalidURN()
+      throws URISyntaxException, RemoteInvocationException {
+    AspectRetriever aspectRetriever = Mockito.mock(AspectRetriever.class);
+    SearchDocumentTransformer searchDocumentTransformer =
+        new SearchDocumentTransformer(1000, 1000, 1000);
+    searchDocumentTransformer.setAspectRetriever(aspectRetriever);
+    EntityRegistry entityRegistry = getTestEntityRegistry();
+    List<Object> urnList = List.of(Urn.createFromString("urn:li:refEntity:1"));
+
+    Mockito.when(aspectRetriever.getEntityRegistry()).thenReturn(entityRegistry);
+    Mockito.when(aspectRetriever.getLatestAspectObject(any(), anyString())).thenReturn(null);
+    SearchableRefFieldSpec searchableRefFieldSpec =
+        entityRegistry.getEntitySpec("testRefEntity").getSearchableRefFieldSpecs().get(0);
+
+    ObjectNode searchDocument = JsonNodeFactory.instance.objectNode();
+    searchDocumentTransformer.setSearchableRefValue(
+        searchableRefFieldSpec, urnList, searchDocument, false);
+    assertTrue(searchDocument.has("refEntityUrns"));
+    assertTrue(searchDocument.get("refEntityUrns").getNodeType().equals(JsonNodeType.NULL));
+  }
+
+  private EntityRegistry getTestEntityRegistry() {
+    return new ConfigEntityRegistry(
+        TestSearchFieldConfig.class
+            .getClassLoader()
+            .getResourceAsStream("test-entity-registry.yaml"));
+  }
 }

metadata-io/src/test/resources/test-entity-registry.yaml

@@ -0,0 +1,10 @@
+id: test-registry
+entities:
+  - name: testRefEntity
+    keyAspect: testRefEntityKey
+    aspects:
+      - testRefEntityInfo
+  - name: refEntity
+    keyAspect: refEntityKey
+    aspects:
+      - refEntityProperties

test-models/src/main/pegasus/com/datahub/test/RefEntityAspect.pdl

@@ -0,0 +1,7 @@
+namespace com.datahub.test
+
+
+/**
+ * A union of all supported metadata aspects for a RefEntity
+ */
+typeref RefEntityAspect = union[RefEntityKey, RefProperties]

test-models/src/main/pegasus/com/datahub/test/RefEntityAssociation.pdl

@@ -0,0 +1,8 @@
+namespace com.datahub.test
+
+import com.linkedin.common.Urn
+import com.linkedin.common.Edge
+
+record RefEntityAssociation includes Edge{
+
+}

test-models/src/main/pegasus/com/datahub/test/RefEntityKey.pdl

@@ -0,0 +1,17 @@
+namespace com.datahub.test
+
+import com.linkedin.common.Urn
+
+/**
+ * Key for Test Entity entity
+ */
+@Aspect = {
+ "name": "refEntityKey"
+}
+record RefEntityKey {
+
+  /**
+    * A unique id
+    */
+    id: string
+}