fix(mae-consumer): fix aspect retriever injections mae-consumer (#10125)

datahub-project · Mar 25, 2024 · f726135 · f726135
1 parent eb2278c
commit f726135
Show file tree

Hide file tree

Showing 2 changed files with 37 additions and 0 deletions.
diff --git a/...ta-jobs/mae-consumer-job/src/main/java/com/linkedin/metadata/kafka/MaeConsumerConfig.java b/...ta-jobs/mae-consumer-job/src/main/java/com/linkedin/metadata/kafka/MaeConsumerConfig.java
@@ -0,0 +1,23 @@
+package com.linkedin.metadata.kafka;
+
+import com.linkedin.metadata.aspect.AspectRetriever;
+import com.linkedin.metadata.search.EntitySearchService;
+import javax.annotation.PostConstruct;
+import org.springframework.beans.factory.annotation.Autowired;
+import org.springframework.beans.factory.annotation.Qualifier;
+import org.springframework.context.annotation.Configuration;
+
+@Configuration
+public class MaeConsumerConfig {
+
+  @Autowired
+  @Qualifier("cachingAspectRetriever")
+  private AspectRetriever aspectRetriever;
+
+  @Autowired private EntitySearchService entitySearchService;
+
+  @PostConstruct
+  protected void postConstruct() {
+    entitySearchService.postConstruct(aspectRetriever);
+  }
+}
diff --git a/...ae-consumer-job/src/test/java/com/linkedin/metadata/kafka/MaeConsumerApplicationTest.java b/...ae-consumer-job/src/test/java/com/linkedin/metadata/kafka/MaeConsumerApplicationTest.java
@@ -2,9 +2,13 @@
 
 import static org.testng.AssertJUnit.*;
 
+import com.linkedin.metadata.aspect.AspectRetriever;
 import com.linkedin.metadata.entity.EntityService;
+import com.linkedin.metadata.search.EntitySearchService;
+import com.linkedin.metadata.search.elasticsearch.query.ESSearchDAO;
 import com.linkedin.metadata.service.FormService;
 import io.datahubproject.metadata.jobs.common.health.kafka.KafkaHealthIndicator;
+import org.apache.commons.lang3.reflect.FieldUtils;
 import org.springframework.beans.factory.annotation.Autowired;
 import org.springframework.boot.test.context.SpringBootTest;
 import org.springframework.test.context.ActiveProfiles;
@@ -22,10 +26,20 @@ public class MaeConsumerApplicationTest extends AbstractTestNGSpringContextTests
 
   @Autowired private FormService formService;
 
+  @Autowired private EntitySearchService entitySearchService;
+
   @Test
   public void testMaeConsumerAutoWiring() {
     assertNotNull(mockEntityService);
     assertNotNull(kafkaHealthIndicator);
     assertNotNull(formService);
   }
+
+  @Test
+  public void testPostConstruct() throws IllegalAccessException {
+    ESSearchDAO test = (ESSearchDAO) FieldUtils.readField(entitySearchService, "esSearchDAO", true);
+    AspectRetriever aspectRetriever =
+        (AspectRetriever) FieldUtils.readField(test, "aspectRetriever", true);
+    assertNotNull(aspectRetriever);
+  }
 }