diff --git a/services/venice-controller/src/main/java/com/linkedin/venice/controller/ZkHelixAdminClient.java b/services/venice-controller/src/main/java/com/linkedin/venice/controller/ZkHelixAdminClient.java index 4726eff9a9..2805a3fc18 100644 --- a/services/venice-controller/src/main/java/com/linkedin/venice/controller/ZkHelixAdminClient.java +++ b/services/venice-controller/src/main/java/com/linkedin/venice/controller/ZkHelixAdminClient.java @@ -103,7 +103,7 @@ public void createVeniceControllerCluster() { helixAdmin.addStateModelDef(controllerClusterName, LeaderStandbySMD.name, LeaderStandbySMD.build()); if (commonConfig.isControllerClusterHelixCloudEnabled()) { - setCloudConfig(commonConfig); + setCloudConfig(controllerClusterName, commonConfig); } } return true; @@ -130,7 +130,7 @@ public void createVeniceStorageCluster(String clusterName, Map h VeniceControllerClusterConfig config = multiClusterConfigs.getControllerConfig(clusterName); if (config.isStorageClusterHelixCloudEnabled()) { - setCloudConfig(config); + setCloudConfig(clusterName, config); } } return true; @@ -327,7 +327,7 @@ public void addInstanceTag(String clusterName, String instanceName, String tag) helixAdmin.addInstanceTag(clusterName, instanceName, tag); } - public void setCloudConfig(VeniceControllerClusterConfig config) { + public void setCloudConfig(String clusterName, VeniceControllerClusterConfig config) { String cloudId = config.getHelixCloudId(); List cloudInfoSources = config.getHelixCloudInfoSources(); String cloudInfoProcessorName = config.getHelixCloudInfoProcessorName(); @@ -345,6 +345,6 @@ public void setCloudConfig(VeniceControllerClusterConfig config) { if (!cloudInfoProcessorName.isEmpty()) { cloudConfigBuilder.setCloudInfoProcessorName(cloudInfoProcessorName); } - helixAdmin.addCloudConfig(controllerClusterName, cloudConfigBuilder.build()); + helixAdmin.addCloudConfig(clusterName, cloudConfigBuilder.build()); } } diff --git a/services/venice-controller/src/test/java/com/linkedin/venice/controller/TestZkHelixAdminClient.java b/services/venice-controller/src/test/java/com/linkedin/venice/controller/TestZkHelixAdminClient.java index 78c3d147d0..d68a652bd5 100644 --- a/services/venice-controller/src/test/java/com/linkedin/venice/controller/TestZkHelixAdminClient.java +++ b/services/venice-controller/src/test/java/com/linkedin/venice/controller/TestZkHelixAdminClient.java @@ -28,6 +28,8 @@ public class TestZkHelixAdminClient { private VeniceControllerMultiClusterConfig mockMultiClusterConfigs; private VeniceControllerClusterConfig mockClusterConfig; + private static final String controllerClusterName = "venice-controllers"; + @BeforeMethod public void setUp() throws NoSuchFieldException, IllegalAccessException { zkHelixAdminClient = mock(ZkHelixAdminClient.class); @@ -78,8 +80,8 @@ public void testSetCloudConfig() { when(mockClusterConfig.getHelixCloudInfoSources()).thenReturn(cloudInfoSources); when(mockClusterConfig.getHelixCloudInfoProcessorName()).thenReturn("TestProcessor"); - doCallRealMethod().when(zkHelixAdminClient).setCloudConfig(mockClusterConfig); - zkHelixAdminClient.setCloudConfig(mockClusterConfig); + doCallRealMethod().when(zkHelixAdminClient).setCloudConfig(controllerClusterName, mockClusterConfig); + zkHelixAdminClient.setCloudConfig(controllerClusterName, mockClusterConfig); verify(mockHelixAdmin).addCloudConfig(any(), any()); } @@ -92,8 +94,10 @@ public void testControllerCloudInfoSourcesNotSet() { when(mockClusterConfig.getHelixCloudInfoSources()).thenReturn(Collections.emptyList()); when(mockClusterConfig.getHelixCloudInfoProcessorName()).thenReturn("TestProcessor"); - doCallRealMethod().when(zkHelixAdminClient).setCloudConfig(mockClusterConfig); - assertThrows(HelixException.class, () -> zkHelixAdminClient.setCloudConfig(mockClusterConfig)); + doCallRealMethod().when(zkHelixAdminClient).setCloudConfig(controllerClusterName, mockClusterConfig); + assertThrows( + HelixException.class, + () -> zkHelixAdminClient.setCloudConfig(controllerClusterName, mockClusterConfig)); } @Test @@ -107,7 +111,9 @@ public void testControllerCloudInfoProcessorNameNotSet() { when(mockClusterConfig.getHelixCloudInfoSources()).thenReturn(cloudInfoSources); when(mockClusterConfig.getHelixCloudInfoProcessorName()).thenReturn(""); - doCallRealMethod().when(zkHelixAdminClient).setCloudConfig(mockClusterConfig); - assertThrows(HelixException.class, () -> zkHelixAdminClient.setCloudConfig(mockClusterConfig)); + doCallRealMethod().when(zkHelixAdminClient).setCloudConfig(controllerClusterName, mockClusterConfig); + assertThrows( + HelixException.class, + () -> zkHelixAdminClient.setCloudConfig(controllerClusterName, mockClusterConfig)); } }