Assert.assertEquals(matches.size(), 2);
Assert.assertTrue(containsConstraint(matches, constraint1));
Assert.assertTrue(containsConstraint(matches, constraint3));
// test messageThrottleStage
ClusterEvent event = new ClusterEvent("testEvent");
event.addAttribute("helixmanager", manager);
Pipeline dataRefresh = new Pipeline();
dataRefresh.addStage(new ReadClusterDataStage());
runPipeline(event, dataRefresh);
runStage(event, new ResourceComputationStage());
MessageSelectionStageOutput msgSelectOutput = new MessageSelectionStageOutput();
Message msg3 =
createMessage(MessageType.STATE_TRANSITION, "msgId-003", "OFFLINE", "SLAVE", "TestDB",
"localhost_0");
Message msg4 =
createMessage(MessageType.STATE_TRANSITION, "msgId-004", "OFFLINE", "SLAVE", "TestDB",
"localhost_0");
Message msg5 =
createMessage(MessageType.STATE_TRANSITION, "msgId-005", "OFFLINE", "SLAVE", "TestDB",
"localhost_0");
Message msg6 =
createMessage(MessageType.STATE_TRANSITION, "msgId-006", "OFFLINE", "SLAVE", "TestDB",
"localhost_1");
List<Message> selectMessages = new ArrayList<Message>();
selectMessages.add(msg1);
selectMessages.add(msg2);
selectMessages.add(msg3);
selectMessages.add(msg4);
selectMessages.add(msg5); // should be throttled
selectMessages.add(msg6); // should be throttled
msgSelectOutput.addMessages("TestDB", new Partition("TestDB_0"), selectMessages);
event.addAttribute(AttributeName.MESSAGES_SELECTED.toString(), msgSelectOutput);
runStage(event, throttleStage);
MessageThrottleStageOutput msgThrottleOutput =
event.getAttribute(AttributeName.MESSAGES_THROTTLE.toString());
List<Message> throttleMessages =
msgThrottleOutput.getMessages("TestDB", new Partition("TestDB_0"));
Assert.assertEquals(throttleMessages.size(), 4);
Assert.assertTrue(throttleMessages.contains(msg1));
Assert.assertTrue(throttleMessages.contains(msg2));