fix(config): fix mcp batch configuration (#13598)

metadata-io/src/test/resources/testng.xml

@@ -13,4 +13,4 @@ parallel followed by everything else.
     <suite-file path="testng-other.xml"/>
     <suite-file path="testng-sql-opt.xml"/>
 </suite-files>
-</suite>
+</suite>

metadata-jobs/mce-consumer/src/main/java/com/linkedin/metadata/kafka/batch/BatchMetadataChangeProposalsProcessor.java

@@ -149,7 +149,7 @@ public class BatchMetadataChangeProposalsProcessor {
       // If adding this MCP would exceed the batch size limit, process the current batch first
       if (!currentBatch.isEmpty()
           && currentBatchSize + mcpSize
-              > provider.getMetadataChangeProposal().getBatch().getSize()) {
+              > provider.getMetadataChangeProposal().getConsumer().getBatch().getSize()) {
         processBatch(currentBatch, currentBatchSize);
         totalProcessed += currentBatch.size();
         log.info(

metadata-jobs/mce-consumer/src/test/java/com/linkedin/metadata/kafka/batch/BatchMetadataChangeProposalsProcessorTest.java

@@ -242,9 +242,11 @@ public class BatchMetadataChangeProposalsProcessorTest {
     when(mockProvider.getMetadataChangeProposal())
         .thenReturn(
             new MetadataChangeProposalConfig()
-                .setBatch(
+                .setConsumer(
                     new MetadataChangeProposalConfig.ConsumerBatchConfig()
-                        .setSize(Integer.MAX_VALUE)));
+                        .setBatch(
+                            new MetadataChangeProposalConfig.BatchConfig()
+                                .setSize(Integer.MAX_VALUE))));
 
     // Create MCPs
     MetadataChangeProposal mcp1 = new MetadataChangeProposal();
@@ -302,9 +304,11 @@ public class BatchMetadataChangeProposalsProcessorTest {
     when(mockProvider.getMetadataChangeProposal())
         .thenReturn(
             new MetadataChangeProposalConfig()
-                .setBatch(
+                .setConsumer(
                     new MetadataChangeProposalConfig.ConsumerBatchConfig()
-                        .setSize(Integer.MAX_VALUE)));
+                        .setBatch(
+                            new MetadataChangeProposalConfig.BatchConfig()
+                                .setSize(Integer.MAX_VALUE))));
 
     // Create 3 Invalid MCPs
     MetadataChangeProposal mcp1 = new MetadataChangeProposal();
@@ -412,11 +416,14 @@ public class BatchMetadataChangeProposalsProcessorTest {
   @Test
   public void testLargeBatchPartitioning() throws Exception {
     // Mock the ConfigurationProvider to return a specific batch size limit
-    MetadataChangeProposalConfig.ConsumerBatchConfig mockConfig =
+    MetadataChangeProposalConfig.ConsumerBatchConfig batchConfig =
-        mock(MetadataChangeProposalConfig.ConsumerBatchConfig.class);
+        new MetadataChangeProposalConfig.ConsumerBatchConfig()
-    when(mockConfig.getSize()).thenReturn(5 * 1024); // 5KB batch size limit for testing
+            .setBatch(
+                new MetadataChangeProposalConfig.BatchConfig()
+                    .setSize(5 * 1024)
+                    .setEnabled(true)); // 5KB batch size limit for testing
     when(mockProvider.getMetadataChangeProposal())
-        .thenReturn(new MetadataChangeProposalConfig().setBatch(mockConfig));
+        .thenReturn(new MetadataChangeProposalConfig().setConsumer(batchConfig));
 
     // Create 3 MCPs, one with a large aspect value
     MetadataChangeProposal smallMcp1 = createMcpWithAspectSize(1000); // 1KB
@@ -456,11 +463,15 @@ public class BatchMetadataChangeProposalsProcessorTest {
   @Test
   public void testExtremelyLargeAspect() throws Exception {
     // Mock the ConfigurationProvider to return a specific batch size limit
-    MetadataChangeProposalConfig.ConsumerBatchConfig mockConfig =
+    MetadataChangeProposalConfig.ConsumerBatchConfig batchConfig =
-        mock(MetadataChangeProposalConfig.ConsumerBatchConfig.class);
+        new MetadataChangeProposalConfig.ConsumerBatchConfig()
-    when(mockConfig.getSize()).thenReturn(10000); // 10KB batch size limit for testing
+            .setBatch(
+                new MetadataChangeProposalConfig.BatchConfig()
+                    .setSize(10000)
+                    .setEnabled(true)); // 10KB batch size limit for testing
     when(mockProvider.getMetadataChangeProposal())
-        .thenReturn(new MetadataChangeProposalConfig().setBatch(mockConfig));
+        .thenReturn(new MetadataChangeProposalConfig().setConsumer(batchConfig));
+    mock(MetadataChangeProposalConfig.ConsumerBatchConfig.class);
 
     // Create an MCP with an aspect value that exceeds the batch size on its own
     MetadataChangeProposal hugeMcp =

metadata-service/configuration/build.gradle

@@ -16,6 +16,9 @@ dependencies {
     compileOnly externalDependency.lombok
 
     annotationProcessor externalDependency.lombok
+
+    testImplementation externalDependency.testng
+    testImplementation externalDependency.springBootTest
 }
 
 processResources.configure {

metadata-service/configuration/src/main/java/com/linkedin/metadata/config/DataHubAppConfiguration.java

@@ -7,8 +7,12 @@ import com.linkedin.metadata.config.search.ElasticSearchConfiguration;
 import com.linkedin.metadata.config.search.SearchServiceConfiguration;
 import com.linkedin.metadata.config.telemetry.TelemetryConfiguration;
 import lombok.Data;
+import org.springframework.boot.context.properties.ConfigurationProperties;
+import org.springframework.stereotype.Component;
 
 @Data
+@Component
+@ConfigurationProperties
 public class DataHubAppConfiguration {
 
   /** Ingestion related configs */

metadata-service/configuration/src/main/java/com/linkedin/metadata/config/MetadataChangeProposalConfig.java

@@ -6,7 +6,7 @@ import lombok.experimental.Accessors;
 @Data
 @Accessors(chain = true)
 public class MetadataChangeProposalConfig {
-  ConsumerBatchConfig batch;
+  ConsumerBatchConfig consumer;
   ThrottlesConfig throttle;
   MCPValidationConfig validation;
   SideEffectsConfig sideEffects;
@@ -64,8 +64,14 @@ public class MetadataChangeProposalConfig {
 
   @Data
   @Accessors(chain = true)
-  public static class ConsumerBatchConfig {
+  public static class BatchConfig {
     boolean enabled;
     Integer size;
   }
+
+  @Data
+  @Accessors(chain = true)
+  public static class ConsumerBatchConfig {
+    BatchConfig batch;
+  }
 }

metadata-service/configuration/src/test/java/com/linkedin/metadata/config/DataHubAppConfigurationTest.java

@@ -0,0 +1,40 @@
+package com.linkedin.metadata.config;
+
+import static org.testng.Assert.assertEquals;
+import static org.testng.Assert.assertFalse;
+import static org.testng.Assert.assertNotNull;
+
+import org.springframework.beans.factory.annotation.Autowired;
+import org.springframework.boot.test.context.SpringBootTest;
+import org.springframework.test.context.testng.AbstractTestNGSpringContextTests;
+import org.testng.annotations.Test;
+
+@SpringBootTest(classes = DataHubTestApplication.class)
+public class DataHubAppConfigurationTest extends AbstractTestNGSpringContextTests {
+
+  @Autowired private DataHubTestApplication testApplication;
+
+  @Test
+  public void testInit() {
+    assertNotNull(testApplication);
+  }
+
+  @Test
+  public void testMCPBatchDefaults() {
+    assertFalse(
+        testApplication
+            .getDataHubAppConfig()
+            .getMetadataChangeProposal()
+            .getConsumer()
+            .getBatch()
+            .isEnabled());
+    assertEquals(
+        testApplication
+            .getDataHubAppConfig()
+            .getMetadataChangeProposal()
+            .getConsumer()
+            .getBatch()
+            .getSize(),
+        15744000);
+  }
+}

metadata-service/configuration/src/test/java/com/linkedin/metadata/config/DataHubTestApplication.java

@@ -0,0 +1,16 @@
+package com.linkedin.metadata.config;
+
+import com.linkedin.metadata.spring.YamlPropertySourceFactory;
+import org.springframework.beans.factory.annotation.Autowired;
+import org.springframework.boot.autoconfigure.SpringBootApplication;
+import org.springframework.context.annotation.PropertySource;
+
+@SpringBootApplication
+@PropertySource(value = "classpath:/application.yaml", factory = YamlPropertySourceFactory.class)
+public class DataHubTestApplication {
+  @Autowired private DataHubAppConfiguration dataHubAppConfig;
+
+  public DataHubAppConfiguration getDataHubAppConfig() {
+    return dataHubAppConfig;
+  }
+}