diff --git a/metadata-jobs/mce-consumer-job/src/test/java/com/linkedin/metadata/kafka/MceConsumerApplicationTestConfiguration.java b/metadata-jobs/mce-consumer-job/src/test/java/com/linkedin/metadata/kafka/MceConsumerApplicationTestConfiguration.java index 4b70d92f6d..747a4ab312 100644 --- a/metadata-jobs/mce-consumer-job/src/test/java/com/linkedin/metadata/kafka/MceConsumerApplicationTestConfiguration.java +++ b/metadata-jobs/mce-consumer-job/src/test/java/com/linkedin/metadata/kafka/MceConsumerApplicationTestConfiguration.java @@ -8,13 +8,13 @@ import com.linkedin.metadata.graph.SiblingGraphService; import com.linkedin.metadata.models.registry.ConfigEntityRegistry; import com.linkedin.metadata.models.registry.EntityRegistry; import com.linkedin.metadata.restli.DefaultRestliClientFactory; -import com.linkedin.metadata.timeseries.elastic.ElasticSearchTimeseriesAspectService; +import com.linkedin.metadata.timeseries.TimeseriesAspectService; import com.linkedin.parseq.retry.backoff.ExponentialBackoff; import com.linkedin.restli.client.Client; import io.ebean.EbeanServer; -import org.mockito.Mockito; import org.springframework.beans.factory.annotation.Autowired; import org.springframework.boot.test.context.TestConfiguration; +import org.springframework.boot.test.mock.mockito.MockBean; import org.springframework.boot.test.web.client.TestRestTemplate; import org.springframework.context.annotation.Bean; import org.springframework.context.annotation.Import; @@ -29,17 +29,11 @@ public class MceConsumerApplicationTestConfiguration { @Autowired private TestRestTemplate restTemplate; - @Bean("kafkaHealthChecker") - @Primary - public KafkaHealthChecker kafkaHealthChecker() { - return Mockito.mock(KafkaHealthChecker.class); - } + @MockBean + public KafkaHealthChecker kafkaHealthChecker; - @Bean("entityService") - @Primary - public EntityService entityService() { - return Mockito.mock(EntityService.class); - } + @MockBean + public EntityService entityService; @Bean("restliEntityClient") @Primary @@ -49,32 +43,18 @@ public class MceConsumerApplicationTestConfiguration { return new RestliEntityClient(restClient, new ExponentialBackoff(1), 1); } - @Bean - @Primary - public EbeanServer ebeanServer() { - return Mockito.mock(EbeanServer.class); - } + @MockBean + public EbeanServer ebeanServer; - @Bean(name = "elasticSearchTimeseriesAspectService") - @Primary - protected ElasticSearchTimeseriesAspectService elasticSearchTimeseriesAspectService() { - return Mockito.mock(ElasticSearchTimeseriesAspectService.class); - } + @MockBean + protected TimeseriesAspectService timeseriesAspectService; - @Bean("entityRegistry") - @Primary - protected EntityRegistry entityRegistry() { - return Mockito.mock(EntityRegistry.class); - } + @MockBean + protected EntityRegistry entityRegistry; - @Bean("configEntityRegistry") - protected ConfigEntityRegistry configEntityRegistry() { - return Mockito.mock(ConfigEntityRegistry.class); - } + @MockBean + protected ConfigEntityRegistry configEntityRegistry; - @Bean("siblingGraphService") - @Primary - protected SiblingGraphService siblingGraphService() { - return Mockito.mock(SiblingGraphService.class); - } + @MockBean + protected SiblingGraphService siblingGraphService; }