Skip to content

Commit

Permalink
[DNM] Adding param for utility
Browse files Browse the repository at this point in the history
  • Loading branch information
@aditya0811
aditya0811 committed Jan 24, 2024
1 parent 9df8594 commit ce298b7
Showing 1 changed file with 35 additions and 36 deletions.
71 changes: 35 additions & 36 deletions ...n/src/test/java/org/apache/pinot/common/assignment/InstanceAssignmentConfigUtilsTest.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -45,7 +45,8 @@ public class InstanceAssignmentConfigUtilsTest {
@Test
public void testShouldRelocateCompletedSegmentsWhenInstancePartitionIsCompleted() {
Map<String, InstanceAssignmentConfig> instanceAssignmentConfigMap = new HashMap<>();
instanceAssignmentConfigMap.put(InstancePartitionsType.COMPLETED.toString(), getInstanceAssignmentConfig());
instanceAssignmentConfigMap.put(InstancePartitionsType.COMPLETED.toString(),
getInstanceAssignmentConfig(InstanceAssignmentConfig.PartitionSelector.FD_AWARE_INSTANCE_PARTITION_SELECTOR));
TableConfig tableConfig = new TableConfigBuilder(TableType.OFFLINE).setTableName("testTable")
.setInstanceAssignmentConfigMap(instanceAssignmentConfigMap).build();
Assert.assertTrue(InstanceAssignmentConfigUtils.shouldRelocateCompletedSegments(tableConfig));
Expand All @@ -55,7 +56,8 @@ public void testShouldRelocateCompletedSegmentsWhenInstancePartitionIsCompleted(
@Test
public void testShouldRelocateCompletedSegmentsWhenInstancePartitionIsConsuming() {
Map<String, InstanceAssignmentConfig> instanceAssignmentConfigMap = new HashMap<>();
instanceAssignmentConfigMap.put(InstancePartitionsType.CONSUMING.name(), getInstanceAssignmentConfig());
instanceAssignmentConfigMap.put(InstancePartitionsType.CONSUMING.name(),
getInstanceAssignmentConfig(InstanceAssignmentConfig.PartitionSelector.FD_AWARE_INSTANCE_PARTITION_SELECTOR));
TagOverrideConfig tagOverrideConfig = new TagOverrideConfig("broker", "Server");
TableConfig tableConfig = new TableConfigBuilder(TableType.OFFLINE).setTableName("testTable")
.setInstanceAssignmentConfigMap(instanceAssignmentConfigMap).setTagOverrideConfig(tagOverrideConfig).build();
Expand All @@ -67,7 +69,8 @@ public void testShouldRelocateCompletedSegmentsWhenInstancePartitionIsConsuming(
@Test
public void testAllowInstanceAssignmentWithPreConfiguredInstancePartitions() {
Map<String, InstanceAssignmentConfig> instanceAssignmentConfigMap = new HashMap<>();
instanceAssignmentConfigMap.put(InstancePartitionsType.CONSUMING.name(), getInstanceAssignmentConfig());
instanceAssignmentConfigMap.put(InstancePartitionsType.CONSUMING.name(),
getInstanceAssignmentConfig(InstanceAssignmentConfig.PartitionSelector.FD_AWARE_INSTANCE_PARTITION_SELECTOR));
Map<InstancePartitionsType, String> instancePartitionsTypeStringMap = new HashMap<>();
instancePartitionsTypeStringMap.put(InstancePartitionsType.CONSUMING, "testTable");
TableConfig tableConfig = new TableConfigBuilder(TableType.OFFLINE).setTableName("testTable")
Expand All @@ -84,7 +87,8 @@ public void testAllowInstanceAssignmentWithPreConfiguredInstancePartitions() {
public void testAllowInstanceAssignmentWithoutPreConfiguredInstancePartitionsOffline() {
Map<String, InstanceAssignmentConfig> instanceAssignmentConfigMap = new HashMap<>();
//check map key
instanceAssignmentConfigMap.put(InstancePartitionsType.OFFLINE.name(), getInstanceAssignmentConfig());
instanceAssignmentConfigMap.put(InstancePartitionsType.OFFLINE.name(),
getInstanceAssignmentConfig(InstanceAssignmentConfig.PartitionSelector.FD_AWARE_INSTANCE_PARTITION_SELECTOR));
Map<InstancePartitionsType, String> instancePartitionsTypeStringMap = new HashMap<>();
instancePartitionsTypeStringMap.put(InstancePartitionsType.CONSUMING, "testTable");
TableConfig tableConfig = new TableConfigBuilder(TableType.OFFLINE).setTableName("testTable")
Expand All @@ -99,7 +103,8 @@ public void testAllowInstanceAssignmentWithoutPreConfiguredInstancePartitionsOff
public void testAllowInstanceAssignmentWithoutPreConfiguredInstancePartitionsConsuming() {
Map<String, InstanceAssignmentConfig> instanceAssignmentConfigMap = new HashMap<>();
//check map key
instanceAssignmentConfigMap.put(InstancePartitionsType.CONSUMING.name(), getInstanceAssignmentConfig());
instanceAssignmentConfigMap.put(InstancePartitionsType.CONSUMING.name(),
getInstanceAssignmentConfig(InstanceAssignmentConfig.PartitionSelector.FD_AWARE_INSTANCE_PARTITION_SELECTOR));

Map<InstancePartitionsType, String> instancePartitionsTypeStringMap = new HashMap<>();
instancePartitionsTypeStringMap.put(InstancePartitionsType.CONSUMING, "testTable");
Expand All @@ -114,7 +119,8 @@ public void testAllowInstanceAssignmentWithoutPreConfiguredInstancePartitionsCon
@Test
public void testAllowInstanceAssignmentWithoutPreConfiguredInstancePartitionsCOMPLETED() {
Map<String, InstanceAssignmentConfig> instanceAssignmentConfigMap = new HashMap<>();
instanceAssignmentConfigMap.put(InstancePartitionsType.COMPLETED.name(), getInstanceAssignmentConfig());
instanceAssignmentConfigMap.put(InstancePartitionsType.COMPLETED.name(),
getInstanceAssignmentConfig(InstanceAssignmentConfig.PartitionSelector.FD_AWARE_INSTANCE_PARTITION_SELECTOR));
Map<InstancePartitionsType, String> instancePartitionsTypeStringMap = new HashMap<>();
instancePartitionsTypeStringMap.put(InstancePartitionsType.COMPLETED, "testTable");
TableConfig tableConfig = new TableConfigBuilder(TableType.REALTIME).setTableName("testTable")
Expand Down Expand Up @@ -154,7 +160,8 @@ public void testAllowInstanceAssignmentWithoutPreConfiguredInstancePartitionsCOM
public void testGetInstanceAssignmentConfigWhenInstanceAssignmentConfig() {
Map<String, InstanceAssignmentConfig> instanceAssignmentConfigMap = new HashMap<>();
//check map key
instanceAssignmentConfigMap.put(InstancePartitionsType.COMPLETED.name(), getInstanceAssignmentConfig());
instanceAssignmentConfigMap.put(InstancePartitionsType.COMPLETED.name(),
getInstanceAssignmentConfig(InstanceAssignmentConfig.PartitionSelector.FD_AWARE_INSTANCE_PARTITION_SELECTOR));

Map<InstancePartitionsType, String> instancePartitionsTypeStringMap = new HashMap<>();
instancePartitionsTypeStringMap.put(InstancePartitionsType.COMPLETED, "testTable");
Expand All @@ -166,9 +173,9 @@ public void testGetInstanceAssignmentConfigWhenInstanceAssignmentConfig() {
.getConstraints().get(0),
"constraints1");
//Check if you can assert entire objects with values in line 214
Assert.assertEquals(InstanceAssignmentConfigUtils.getInstanceAssignmentConfig(tableConfig, InstancePartitionsType.COMPLETED)
.getReplicaGroupPartitionConfig().getNumInstancesPerPartition(), 10);

Assert.assertEquals(InstanceAssignmentConfigUtils.getInstanceAssignmentConfig(tableConfig,
InstancePartitionsType.COMPLETED).getReplicaGroupPartitionConfig().getNumInstancesPerPartition(),
1);
}

@Test
Expand All @@ -195,12 +202,12 @@ public void testGetInstanceAssignmentConfigWhenInstanceAssignmentConfigIsNotPres
SegmentPartitionConfig segmentPartitionConfig = new SegmentPartitionConfig(columnPartitionConfigMap);
indexingConfig.setSegmentPartitionConfig(segmentPartitionConfig);
tableConfig.setIndexingConfig(indexingConfig);
Assert.assertEquals(InstanceAssignmentConfigUtils.getInstanceAssignmentConfig(tableConfig, InstancePartitionsType.OFFLINE)
.getReplicaGroupPartitionConfig().isReplicaGroupBased(), Boolean.TRUE);
Assert.assertEquals(InstanceAssignmentConfigUtils.getInstanceAssignmentConfig(tableConfig, InstancePartitionsType.OFFLINE)
.getReplicaGroupPartitionConfig().getPartitionColumn(), "column1");
Assert.assertEquals(InstanceAssignmentConfigUtils.getInstanceAssignmentConfig(tableConfig, InstancePartitionsType.OFFLINE)
.getReplicaGroupPartitionConfig().getNumInstancesPerPartition(), 1);
Assert.assertEquals(InstanceAssignmentConfigUtils.getInstanceAssignmentConfig(tableConfig,
InstancePartitionsType.OFFLINE).getReplicaGroupPartitionConfig().isReplicaGroupBased(), Boolean.TRUE);
Assert.assertEquals(InstanceAssignmentConfigUtils.getInstanceAssignmentConfig(tableConfig,
InstancePartitionsType.OFFLINE).getReplicaGroupPartitionConfig().getPartitionColumn(), "column1");
Assert.assertEquals(InstanceAssignmentConfigUtils.getInstanceAssignmentConfig(tableConfig,
InstancePartitionsType.OFFLINE).getReplicaGroupPartitionConfig().getNumInstancesPerPartition(), 1);
//Check if we want to asset all the objects for getReplicaGroupPartitionConfig()
}

Expand Down Expand Up @@ -247,40 +254,32 @@ public void testGetInstanceAssignmentConfigWhenInstanceAssignmentConfigIsNotPres
public void testIsMirrorServerSetAssignment() {

Map<String, InstanceAssignmentConfig> instanceAssignmentConfigMap = new HashMap<>();
InstanceTagPoolConfig instanceTagPoolConfig = new InstanceTagPoolConfig("tag",true,1, null);
List<String> constraints = new ArrayList<>();
constraints.add("constraints1");
InstanceConstraintConfig instanceConstraintConfig = new InstanceConstraintConfig(constraints);
InstanceReplicaGroupPartitionConfig instanceReplicaGroupPartitionConfig =
new InstanceReplicaGroupPartitionConfig(true, 1,1,1,1,1,true,null);

InstanceAssignmentConfig instanceAssignmentConfig = new InstanceAssignmentConfig(instanceTagPoolConfig, instanceConstraintConfig, instanceReplicaGroupPartitionConfig, InstanceAssignmentConfig.PartitionSelector.MIRROR_SERVER_SET_PARTITION_SELECTOR.name());
instanceAssignmentConfigMap.put(InstancePartitionsType.OFFLINE.name(), instanceAssignmentConfig);
instanceAssignmentConfigMap.put(InstancePartitionsType.OFFLINE.name(),
getInstanceAssignmentConfig(InstanceAssignmentConfig.PartitionSelector.MIRROR_SERVER_SET_PARTITION_SELECTOR));
TableConfig tableConfig = new TableConfigBuilder(TableType.OFFLINE).setTableName("testTable")
.setInstanceAssignmentConfigMap(instanceAssignmentConfigMap)
.build();

Assert.assertTrue(InstanceAssignmentConfigUtils.isMirrorServerSetAssignment(tableConfig, InstancePartitionsType.OFFLINE));

Assert.assertTrue(InstanceAssignmentConfigUtils.isMirrorServerSetAssignment(tableConfig,
InstancePartitionsType.OFFLINE));
}

//check if you can get False for this method.

private static InstanceAssignmentConfig getInstanceAssignmentConfig() {
private static InstanceAssignmentConfig getInstanceAssignmentConfig(InstanceAssignmentConfig.PartitionSelector
partitionSelector) {
InstanceTagPoolConfig instanceTagPoolConfig =
new InstanceTagPoolConfig("tag",true,1, null);
List<String> constraints = new ArrayList<>();
constraints.add("constraints1");
InstanceConstraintConfig instanceConstraintConfig = new InstanceConstraintConfig(constraints);
InstanceReplicaGroupPartitionConfig instanceReplicaGroupPartitionConfig =
new InstanceReplicaGroupPartitionConfig(true, 1,1,1,1,1,true,null);
InstanceAssignmentConfig instanceAssignmentConfig = new InstanceAssignmentConfig(instanceTagPoolConfig,instanceConstraintConfig,
instanceReplicaGroupPartitionConfig, InstanceAssignmentConfig.PartitionSelector.FD_AWARE_INSTANCE_PARTITION_SELECTOR.name());
new InstanceReplicaGroupPartitionConfig(true, 1, 1,
1, 1, 1, true,
null);
InstanceAssignmentConfig instanceAssignmentConfig
= new InstanceAssignmentConfig(instanceTagPoolConfig, instanceConstraintConfig,
instanceReplicaGroupPartitionConfig, partitionSelector.name());
return instanceAssignmentConfig;
}





}

0 comments on commit ce298b7

Please sign in to comment.