test(mce-consumer): mockbeans (datahub-project#6878)

CybercentreCanada · Feb 1, 2023 · 86ddbb7 · 86ddbb7
1 parent e7cc919
commit 86ddbb7
Showing 1 changed file with 16 additions and 36 deletions.
diff --git a/...ob/src/test/java/com/linkedin/metadata/kafka/MceConsumerApplicationTestConfiguration.java b/...ob/src/test/java/com/linkedin/metadata/kafka/MceConsumerApplicationTestConfiguration.java
@@ -8,13 +8,13 @@
 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 RestliEntityClient restliEntityClient() {
         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;
 }