Java Code Examples for org.apache.helix.model.Message#setMsgId()

private void syncSessionToController(HelixManager manager) {
  if (_lastSessionSyncTime == null ||
          System.currentTimeMillis() - _lastSessionSyncTime > SESSION_SYNC_INTERVAL) { // > delay since last sync
    HelixDataAccessor accessor = manager.getHelixDataAccessor();
    PropertyKey key = new Builder(manager.getClusterName()).controllerMessage(SESSION_SYNC);
    if (accessor.getProperty(key) == null) {
      LOG.info(String.format("Participant %s syncs session with controller", manager.getInstanceName()));
      Message msg = new Message(MessageType.PARTICIPANT_SESSION_CHANGE, SESSION_SYNC);
      msg.setSrcName(manager.getInstanceName());
      msg.setTgtSessionId("*");
      msg.setMsgState(MessageState.NEW);
      msg.setMsgId(SESSION_SYNC);

      Criteria cr = new Criteria();
      cr.setRecipientInstanceType(InstanceType.CONTROLLER);
      cr.setSessionSpecific(false);

      manager.getMessagingService().send(cr, msg);
      _lastSessionSyncTime = System.currentTimeMillis();
    }
  }
}
 

void sendMessage(String msgId, String instanceName, String fromState, String toState,
    String partitionKey, int partitionId) throws InterruptedException, JsonGenerationException,
    JsonMappingException, IOException {
  Message message = new Message(MessageType.STATE_TRANSITION, msgId);
  message.setMsgId(msgId);
  message.setSrcName(srcName);
  message.setTgtName(instanceName);
  message.setMsgState(MessageState.NEW);
  message.setFromState(fromState);
  message.setToState(toState);
  // message.setPartitionId(partitionId);
  message.setPartitionName(partitionKey);

  String path = PropertyPathBuilder.instanceMessage(clusterName, instanceName, message.getId());
  ObjectMapper mapper = new ObjectMapper();
  StringWriter sw = new StringWriter();
  mapper.writeValueUsingView(sw, message, Message.class);
  System.out.println(sw.toString());
  client.delete(path);

  Thread.sleep(10000);
  ZNRecord record = client.readData(PropertyPathBuilder.liveInstance(clusterName, instanceName));
  message.setTgtSessionId(record.getSimpleField(LiveInstanceProperty.SESSION_ID.toString())
      .toString());
  client.createPersistent(path, message);
}
 

private List<Message> generateMessagesForController(Message message) {
  List<Message> messages = new ArrayList<Message>();
  String id = UUID.randomUUID().toString();
  Message newMessage = new Message(message.getRecord(), id);
  newMessage.setMsgId(id);
  newMessage.setSrcName(_manager.getInstanceName());
  newMessage.setTgtName("Controller");
  messages.add(newMessage);
  return messages;
}
 

private List<Message> generateMessagesForController(Message message) {
  List<Message> messages = new ArrayList<Message>();
  String id = (message.getMsgId() == null) ? UUID.randomUUID().toString() : message.getMsgId();
  Message newMessage = new Message(message.getRecord(), id);
  newMessage.setMsgId(id);
  newMessage.setSrcName(_manager.getInstanceName());
  newMessage.setTgtName(InstanceType.CONTROLLER.name());
  messages.add(newMessage);
  return messages;
}
 

public void postTestMessage(String zkServer, String clusterName, String instanceName) {
  String msgSrc = "cm-instance-0";
  String msgId = "TestMessageId-2";

  Message message = new Message(MessageType.STATE_TRANSITION, msgId);
  message.setMsgId(msgId);
  message.setSrcName(msgSrc);
  message.setTgtName(instanceName);
  message.setMsgState(MessageState.NEW);
  message.setFromState("Slave");
  message.setToState("Master");
  message.setPartitionName("EspressoDB.partition-0." + instanceName);

  post(zkServer, message, clusterName, instanceName);
}
 

@Test()
public void TestBlockingSendReceive() throws Exception {
  String hostSrc = "localhost_" + START_PORT;
  String hostDest = "localhost_" + (START_PORT + 1);

  TestMessagingHandlerFactory factory = new TestMessagingHandlerFactory();
  _participants[1].getMessagingService().registerMessageHandlerFactory(factory.getMessageTypes(),
      factory);

  String msgId = new UUID(123, 456).toString();
  Message msg = new Message(factory.getMessageTypes().get(0), msgId);
  msg.setMsgId(msgId);
  msg.setSrcName(hostSrc);

  msg.setTgtSessionId("*");
  msg.setMsgState(MessageState.NEW);
  String para = "Testing messaging para";
  msg.getRecord().setSimpleField("TestMessagingPara", para);

  Criteria cr = new Criteria();
  cr.setInstanceName(hostDest);
  cr.setRecipientInstanceType(InstanceType.PARTICIPANT);
  cr.setSessionSpecific(false);

  AsyncCallback asyncCallback = new MockAsyncCallback();
  int messagesSent =
      _participants[0].getMessagingService().sendAndWait(cr, msg, asyncCallback, 60000);

  AssertJUnit.assertTrue(asyncCallback.getMessageReplied().get(0).getRecord()
      .getMapField(Message.Attributes.MESSAGE_RESULT.toString()).get("ReplyMessage")
      .equals("TestReplyMessage"));
  AssertJUnit.assertTrue(asyncCallback.getMessageReplied().size() == 1);

  AsyncCallback asyncCallback2 = new MockAsyncCallback();
  messagesSent = _participants[0].getMessagingService().sendAndWait(cr, msg, asyncCallback2, 500);
  AssertJUnit.assertTrue(asyncCallback2.isTimedOut());

}
 

@Test()
public void sendSelfMsg() {
  String hostSrc = "localhost_" + START_PORT;

  for (int i = 0; i < NODE_NR; i++) {
    TestMessagingHandlerFactory factory = new TestMessagingHandlerFactory();
    String hostDest = "localhost_" + (START_PORT + i);
    _participants[i].getMessagingService().registerMessageHandlerFactory(
        factory.getMessageTypes(), factory);

  }
  String msgId = new UUID(123, 456).toString();
  Message msg = new Message(new TestMessagingHandlerFactory().getMessageTypes().get(0), msgId);
  msg.setMsgId(msgId);
  msg.setSrcName(hostSrc);

  msg.setTgtSessionId("*");
  msg.setMsgState(MessageState.NEW);
  String para = "Testing messaging para";
  msg.getRecord().setSimpleField("TestMessagingPara", para);

  Criteria cr = new Criteria();
  cr.setInstanceName("%");
  cr.setRecipientInstanceType(InstanceType.PARTICIPANT);
  cr.setSessionSpecific(false);
  cr.setSelfExcluded(false);
  AsyncCallback callback1 = new MockAsyncCallback();
  int messageSent1 =
      _participants[0].getMessagingService().sendAndWait(cr, msg, callback1, 10000);

  AssertJUnit.assertTrue(callback1.getMessageReplied().size() == NODE_NR);
  AssertJUnit.assertTrue(callback1.getMessageReplied().get(0).getRecord()
      .getMapField(Message.Attributes.MESSAGE_RESULT.toString()).get("ReplyMessage")
      .equals("TestReplyMessage"));
}
 

@Test()
public void TestBlockingSendReceive() {
  String hostDest = "localhost_" + (START_PORT + 1);

  TestMessagingHandlerFactory factory = new TestMessagingHandlerFactory();
  _participants[1].getMessagingService().registerMessageHandlerFactory(factory.getMessageTypes(),
      factory);

  String msgId = new UUID(123, 456).toString();
  Message msg = new Message(factory.getMessageTypes().get(0), msgId);
  msg.setMsgId(msgId);
  msg.setSrcName(_hostSrc);
  msg.setTgtSessionId("*");
  msg.setMsgState(MessageState.NEW);
  String para = "Testing messaging para";
  msg.getRecord().setSimpleField("TestMessagingPara", para);

  Criteria cr = new Criteria();
  cr.setInstanceName(hostDest);
  cr.setRecipientInstanceType(InstanceType.PARTICIPANT);
  cr.setSessionSpecific(false);
  cr.setClusterName(CLUSTER_NAME);

  AsyncCallback asyncCallback = new MockAsyncCallback();
  int messagesSent =
      _adminController.getMessagingService().sendAndWait(cr, msg, asyncCallback, 60000);

  AssertJUnit.assertEquals(
      asyncCallback.getMessageReplied().get(0).getRecord()
          .getMapField(Message.Attributes.MESSAGE_RESULT.toString()).get("ReplyMessage"),
      "TestReplyMessage");
  AssertJUnit.assertEquals(asyncCallback.getMessageReplied().size(), 1);

  AsyncCallback asyncCallback2 = new MockAsyncCallback();
  messagesSent = _adminController.getMessagingService().sendAndWait(cr, msg, asyncCallback2, 500);
  AssertJUnit.assertTrue(asyncCallback2.isTimedOut());
}
 

@Test()
public void testInvocation() throws Exception {
  HelixTaskExecutor executor = new HelixTaskExecutor();
  System.out.println("START TestCMTaskHandler.testInvocation()");
  Message message = new Message(MessageType.STATE_TRANSITION, "Some unique id");

  message.setSrcName("cm-instance-0");
  message.setTgtSessionId("1234");
  message.setFromState("Offline");
  message.setToState("Slave");
  message.setPartitionName("TestDB_0");
  message.setMsgId("Some unique message id");
  message.setResourceName("TestDB");
  message.setTgtName("localhost");
  message.setStateModelDef("MasterSlave");
  message.setStateModelFactoryName(HelixConstants.DEFAULT_STATE_MODEL_FACTORY);
  MockMasterSlaveStateModel stateModel = new MockMasterSlaveStateModel();
  NotificationContext context;
  MockManager manager = new MockManager("clusterName");
  HelixDataAccessor accessor = manager.getHelixDataAccessor();
  StateModelDefinition stateModelDef =
      new StateModelDefinition(StateModelConfigGenerator.generateConfigForMasterSlave());
  Builder keyBuilder = accessor.keyBuilder();
  accessor.setProperty(keyBuilder.stateModelDef("MasterSlave"), stateModelDef);

  context = new NotificationContext(manager);
  CurrentState currentStateDelta = new CurrentState("TestDB");
  currentStateDelta.setState("TestDB_0", "OFFLINE");

  HelixStateTransitionHandler stHandler =
      new HelixStateTransitionHandler(null, stateModel, message, context, currentStateDelta);
  HelixTask handler;
  handler = new HelixTask(message, context, stHandler, executor);
  handler.call();
  AssertJUnit.assertTrue(stateModel.stateModelInvoked);
  System.out.println("END TestCMTaskHandler.testInvocation() at "
      + new Date(System.currentTimeMillis()));
}
 

@Test
public void testSyncSessionToController() throws Exception {
  System.out
      .println("START testSyncSessionToController at " + new Date(System.currentTimeMillis()));

  String clusterName = getShortClassName();
  MockParticipantManager[] participants = new MockParticipantManager[5];
  int resourceNb = 10;
  TestHelper.setupCluster(clusterName, ZK_ADDR, 12918, // participant port
      "localhost", // participant name prefix
      "TestDB", // resource name prefix
      resourceNb, // resources
      1, // partitions per resource
      5, // number of nodes
      1, // replicas
      "MasterSlave", true); // do rebalance

  ClusterControllerManager controller =
      new ClusterControllerManager(ZK_ADDR, clusterName, "controller_0");
  controller.syncStart();

  // start participants
  for (int i = 0; i < 5; i++) {
    String instanceName = "localhost_" + (12918 + i);
    participants[i] = new MockParticipantManager(ZK_ADDR, clusterName, instanceName);
    participants[i].syncStart();
  }

  ZKHelixManager zkHelixManager = new ZKHelixManager(clusterName, "controllerMessageListener",
      InstanceType.CONTROLLER, ZK_ADDR);
  zkHelixManager.connect();
  MockMessageListener mockMessageListener = new MockMessageListener();
  zkHelixManager.addControllerMessageListener(mockMessageListener);

  PropertyKey.Builder keyBuilder = new PropertyKey.Builder(clusterName);
  ZkBaseDataAccessor<ZNRecord> accessor = new ZkBaseDataAccessor<>(_gZkClient);
  String path = keyBuilder.liveInstance("localhost_12918").getPath();
  Stat stat = new Stat();
  ZNRecord data = accessor.get(path, stat, 2);
  data.getSimpleFields().put("SESSION_ID", "invalid-id");
  accessor.set(path, data, 2);
  Thread.sleep(2000);
  // Since we always read the content from ephemeral nodes, sync message won't be sent
  Assert.assertFalse(mockMessageListener.isSessionSyncMessageSent());

  // Even after reconnect, session sync won't happen
  ZkTestHelper.expireSession(participants[0].getZkClient());
  Assert.assertFalse(mockMessageListener.isSessionSyncMessageSent());

  // Inject an invalid session message to trigger sync message
  PropertyKey messageKey = keyBuilder.message("localhost_12918", "Mocked Invalid Message");
  Message msg = new Message(Message.MessageType.STATE_TRANSITION, "Mocked Invalid Message");
  msg.setSrcName(controller.getInstanceName());
  msg.setTgtSessionId("invalid-id");
  msg.setMsgState(Message.MessageState.NEW);
  msg.setMsgId("Mocked Invalid Message");
  msg.setTgtName("localhost_12918");
  msg.setPartitionName("foo");
  msg.setResourceName("bar");
  msg.setFromState("SLAVE");
  msg.setToState("MASTER");
  msg.setSrcSessionId(controller.getSessionId());
  msg.setStateModelDef("MasterSlave");
  msg.setStateModelFactoryName("DEFAULT");
  HelixDataAccessor dataAccessor = new ZKHelixDataAccessor(clusterName, accessor);
  dataAccessor.setProperty(messageKey, msg);
  Assert.assertTrue(TestHelper.verify(() -> mockMessageListener.isSessionSyncMessageSent(), 1500));

  // Cleanup
  controller.syncStop();
  zkHelixManager.disconnect();
  for (int i = 0; i < 5; i++) {
    participants[i].syncStop();
  }
  deleteCluster(clusterName);
}
 

@Test()
public void TestMultiMessageCriteria() throws Exception {
  String hostSrc = "localhost_" + START_PORT;

  for (int i = 0; i < NODE_NR; i++) {
    TestMessagingHandlerFactory factory = new TestMessagingHandlerFactory();
    String hostDest = "localhost_" + (START_PORT + i);
    _participants[i].getMessagingService().registerMessageHandlerFactory(
        factory.getMessageTypes(), factory);

  }
  String msgId = new UUID(123, 456).toString();
  Message msg = new Message(new TestMessagingHandlerFactory().getMessageTypes().get(0), msgId);
  msg.setMsgId(msgId);
  msg.setSrcName(hostSrc);

  msg.setTgtSessionId("*");
  msg.setMsgState(MessageState.NEW);
  String para = "Testing messaging para";
  msg.getRecord().setSimpleField("TestMessagingPara", para);

  Criteria cr = new Criteria();
  cr.setInstanceName("%");
  cr.setRecipientInstanceType(InstanceType.PARTICIPANT);
  cr.setSessionSpecific(false);
  AsyncCallback callback1 = new MockAsyncCallback();
  int messageSent1 =
      _participants[0].getMessagingService().sendAndWait(cr, msg, callback1, 10000);

  AssertJUnit.assertTrue(callback1.getMessageReplied().get(0).getRecord()
      .getMapField(Message.Attributes.MESSAGE_RESULT.toString()).get("ReplyMessage")
      .equals("TestReplyMessage"));
  AssertJUnit.assertTrue(callback1.getMessageReplied().size() == NODE_NR - 1);

  AsyncCallback callback2 = new MockAsyncCallback();
  int messageSent2 = _participants[0].getMessagingService().sendAndWait(cr, msg, callback2, 500);

  AssertJUnit.assertTrue(callback2.isTimedOut());

  cr.setPartition("TestDB_17");
  AsyncCallback callback3 = new MockAsyncCallback();
  int messageSent3 =
      _participants[0].getMessagingService().sendAndWait(cr, msg, callback3, 10000);
  AssertJUnit.assertTrue(callback3.getMessageReplied().size() == _replica - 1);

  cr.setPartition("TestDB_15");
  AsyncCallback callback4 = new MockAsyncCallback();
  int messageSent4 =
      _participants[0].getMessagingService().sendAndWait(cr, msg, callback4, 10000);
  AssertJUnit.assertTrue(callback4.getMessageReplied().size() == _replica);

  cr.setPartitionState("SLAVE");
  AsyncCallback callback5 = new MockAsyncCallback();
  int messageSent5 =
      _participants[0].getMessagingService().sendAndWait(cr, msg, callback5, 10000);
  AssertJUnit.assertTrue(callback5.getMessageReplied().size() == _replica - 1);

  cr.setDataSource(DataSource.IDEALSTATES);
  AsyncCallback callback6 = new MockAsyncCallback();
  int messageSent6 =
      _participants[0].getMessagingService().sendAndWait(cr, msg, callback6, 10000);
  AssertJUnit.assertTrue(callback6.getMessageReplied().size() == _replica - 1);
}
 

@Test()
public void TestControllerMessage() throws Exception {
  String hostSrc = "localhost_" + START_PORT;

  for (int i = 0; i < NODE_NR; i++) {
    TestMessagingHandlerFactory factory = new TestMessagingHandlerFactory();
    String hostDest = "localhost_" + (START_PORT + i);
    _participants[i].getMessagingService().registerMessageHandlerFactory(
        factory.getMessageTypes(), factory);

  }
  String msgId = new UUID(123, 456).toString();
  Message msg = new Message(MessageType.CONTROLLER_MSG, msgId);
  msg.setMsgId(msgId);
  msg.setSrcName(hostSrc);

  msg.setTgtSessionId("*");
  msg.setMsgState(MessageState.NEW);
  String para = "Testing messaging para";
  msg.getRecord().setSimpleField("TestMessagingPara", para);

  Criteria cr = new Criteria();
  cr.setInstanceName("*");
  cr.setRecipientInstanceType(InstanceType.CONTROLLER);
  cr.setSessionSpecific(false);

  AsyncCallback callback1 = new MockAsyncCallback();
  int messagesSent =
      _participants[0].getMessagingService().sendAndWait(cr, msg, callback1, 10000);

  AssertJUnit.assertTrue(callback1.getMessageReplied().get(0).getRecord()
      .getMapField(Message.Attributes.MESSAGE_RESULT.toString()).get("ControllerResult")
      .indexOf(hostSrc) != -1);
  AssertJUnit.assertTrue(callback1.getMessageReplied().size() == 1);

  msgId = UUID.randomUUID().toString();
  msg.setMsgId(msgId);
  cr.setPartition("TestDB_17");
  AsyncCallback callback2 = new MockAsyncCallback();
  messagesSent = _participants[0].getMessagingService().sendAndWait(cr, msg, callback2, 10000);

  AssertJUnit.assertTrue(callback2.getMessageReplied().get(0).getRecord()
      .getMapField(Message.Attributes.MESSAGE_RESULT.toString()).get("ControllerResult")
      .indexOf(hostSrc) != -1);

  AssertJUnit.assertTrue(callback2.getMessageReplied().size() == 1);

  msgId = UUID.randomUUID().toString();
  msg.setMsgId(msgId);
  cr.setPartitionState("SLAVE");
  AsyncCallback callback3 = new MockAsyncCallback();
  messagesSent = _participants[0].getMessagingService().sendAndWait(cr, msg, callback3, 10000);
  AssertJUnit.assertTrue(callback3.getMessageReplied().get(0).getRecord()
      .getMapField(Message.Attributes.MESSAGE_RESULT.toString()).get("ControllerResult")
      .indexOf(hostSrc) != -1);

  AssertJUnit.assertTrue(callback3.getMessageReplied().size() == 1);
}
 

@Test()
public void TestMultiMessageCriteria() {
  for (int i = 0; i < NODE_NR; i++) {
    TestMessagingHandlerFactory factory = new TestMessagingHandlerFactory();
    _participants[i].getMessagingService()
        .registerMessageHandlerFactory(factory.getMessageTypes(), factory);
  }
  String msgId = new UUID(123, 456).toString();
  Message msg = new Message(new TestMessagingHandlerFactory().getMessageTypes().get(0), msgId);
  msg.setMsgId(msgId);
  msg.setSrcName(_hostSrc);
  msg.setTgtSessionId("*");
  msg.setMsgState(MessageState.NEW);
  String para = "Testing messaging para";
  msg.getRecord().setSimpleField("TestMessagingPara", para);

  Criteria cr = new Criteria();
  cr.setInstanceName("%");
  cr.setRecipientInstanceType(InstanceType.PARTICIPANT);
  cr.setSessionSpecific(false);
  cr.setClusterName(CLUSTER_NAME);

  AsyncCallback callback1 = new MockAsyncCallback();
  int messageSent1 =
      _adminController.getMessagingService().sendAndWait(cr, msg, callback1, 10000);

  AssertJUnit.assertEquals(
      callback1.getMessageReplied().get(0).getRecord()
          .getMapField(Message.Attributes.MESSAGE_RESULT.toString()).get("ReplyMessage"),
      "TestReplyMessage");
  AssertJUnit.assertEquals(NODE_NR, callback1.getMessageReplied().size());

  AsyncCallback callback2 = new MockAsyncCallback();
  int messageSent2 = _adminController.getMessagingService().sendAndWait(cr, msg, callback2, 500);

  AssertJUnit.assertTrue(callback2.isTimedOut());

  cr.setPartition("TestDB_17");
  AsyncCallback callback3 = new MockAsyncCallback();
  int messageSent3 =
      _adminController.getMessagingService().sendAndWait(cr, msg, callback3, 10000);
  AssertJUnit.assertEquals(_replica, callback3.getMessageReplied().size());

  cr.setPartition("TestDB_15");
  AsyncCallback callback4 = new MockAsyncCallback();
  int messageSent4 =
      _adminController.getMessagingService().sendAndWait(cr, msg, callback4, 10000);
  AssertJUnit.assertEquals(_replica, callback4.getMessageReplied().size());

  cr.setPartitionState("SLAVE");
  AsyncCallback callback5 = new MockAsyncCallback();
  int messageSent5 =
      _adminController.getMessagingService().sendAndWait(cr, msg, callback5, 10000);
  AssertJUnit.assertEquals(_replica - 1, callback5.getMessageReplied().size());

  cr.setDataSource(DataSource.IDEALSTATES);
  AsyncCallback callback6 = new MockAsyncCallback();
  int messageSent6 =
      _adminController.getMessagingService().sendAndWait(cr, msg, callback6, 10000);
  AssertJUnit.assertEquals(_replica - 1, callback6.getMessageReplied().size());
}
 

@Test()
public void TestControllerMessage() {
  for (int i = 0; i < NODE_NR; i++) {
    TestMessagingHandlerFactory factory = new TestMessagingHandlerFactory();
    _participants[i].getMessagingService()
        .registerMessageHandlerFactory(factory.getMessageTypes(), factory);

  }
  String msgId = new UUID(123, 456).toString();
  Message msg = new Message(MessageType.CONTROLLER_MSG, msgId);
  msg.setMsgId(msgId);
  msg.setSrcName(_hostSrc);
  msg.setTgtSessionId("*");
  msg.setMsgState(MessageState.NEW);
  String para = "Testing messaging para";
  msg.getRecord().setSimpleField("TestMessagingPara", para);

  Criteria cr = new Criteria();
  cr.setInstanceName("*");
  cr.setRecipientInstanceType(InstanceType.CONTROLLER);
  cr.setSessionSpecific(false);
  cr.setClusterName(CLUSTER_NAME);

  AsyncCallback callback1 = new MockAsyncCallback();
  int messagesSent =
      _adminController.getMessagingService().sendAndWait(cr, msg, callback1, 10000);

  AssertJUnit.assertTrue(callback1.getMessageReplied().get(0).getRecord()
      .getMapField(Message.Attributes.MESSAGE_RESULT.toString()).get("ControllerResult")
      .contains(_hostSrc));
  AssertJUnit.assertEquals(callback1.getMessageReplied().size(), 1);

  msgId = UUID.randomUUID().toString();
  msg.setMsgId(msgId);
  cr.setPartition("TestDB_17");
  AsyncCallback callback2 = new MockAsyncCallback();
  messagesSent = _adminController.getMessagingService().sendAndWait(cr, msg, callback2, 10000);

  AssertJUnit.assertTrue(callback2.getMessageReplied().get(0).getRecord()
      .getMapField(Message.Attributes.MESSAGE_RESULT.toString()).get("ControllerResult")
      .contains(_hostSrc));

  AssertJUnit.assertEquals(callback2.getMessageReplied().size(), 1);

  msgId = UUID.randomUUID().toString();
  msg.setMsgId(msgId);
  cr.setPartitionState("SLAVE");
  AsyncCallback callback3 = new MockAsyncCallback();
  messagesSent = _adminController.getMessagingService().sendAndWait(cr, msg, callback3, 10000);
  AssertJUnit.assertTrue(callback3.getMessageReplied().get(0).getRecord()
      .getMapField(Message.Attributes.MESSAGE_RESULT.toString()).get("ControllerResult")
      .contains(_hostSrc));

  AssertJUnit.assertEquals(callback3.getMessageReplied().size(), 1);
}
 

@Test()
public void testCMTaskExecutor() throws Exception {
  System.out.println("START TestCMTaskExecutor");
  String msgId = "TestMessageId";
  Message message = new Message(MessageType.TASK_REPLY, msgId);

  message.setMsgId(msgId);
  message.setSrcName("cm-instance-0");
  message.setTgtName("cm-instance-1");
  message.setTgtSessionId("1234");
  message.setFromState("Offline");
  message.setToState("Slave");
  message.setPartitionName("TestDB_0");
  message.setResourceName("TestDB");
  message.setStateModelDef("MasterSlave");

  MockManager manager = new MockManager("clusterName");
  HelixDataAccessor accessor = manager.getHelixDataAccessor();
  StateModelDefinition stateModelDef =
      new StateModelDefinition(StateModelConfigGenerator.generateConfigForMasterSlave());
  Builder keyBuilder = accessor.keyBuilder();
  accessor.setProperty(keyBuilder.stateModelDef("MasterSlave"), stateModelDef);

  MockHelixTaskExecutor executor = new MockHelixTaskExecutor();
  MockMasterSlaveStateModel stateModel = new MockMasterSlaveStateModel();
  executor.registerMessageHandlerFactory(MessageType.TASK_REPLY.name(),
      new AsyncCallbackService());

  NotificationContext context = new NotificationContext(manager);
  CurrentState currentStateDelta = new CurrentState("TestDB");
  currentStateDelta.setState("TestDB_0", "OFFLINE");

  StateModelFactory<MockMasterSlaveStateModel> stateModelFactory = new StateModelFactory<MockMasterSlaveStateModel>() {

    @Override
    public MockMasterSlaveStateModel createNewStateModel(String resource, String partitionName) {
      // TODO Auto-generated method stub
      return new MockMasterSlaveStateModel();
    }

  };
  HelixStateTransitionHandler handler =
      new HelixStateTransitionHandler(stateModelFactory, stateModel, message, context,
          currentStateDelta);

  HelixTask task = new HelixTask(message, context, handler, executor);
  executor.scheduleTask(task);
  for (int i = 0; i < 10; i++) {
    if (!executor.isDone(task.getTaskId())) {
      Thread.sleep(500);
    }
  }
  AssertJUnit.assertTrue(stateModel.stateModelInvoked);
  System.out.println("END TestCMTaskExecutor");
}
 

@Test()
public void testInvocationAnnotated() throws Exception {
  System.out.println("START TestCMTaskHandler.testInvocationAnnotated() at "
      + new Date(System.currentTimeMillis()));
  HelixTaskExecutor executor = new HelixTaskExecutor();
  Message message = new Message(MessageType.STATE_TRANSITION, "Some unique id");
  message.setSrcName("cm-instance-0");
  message.setTgtSessionId("1234");
  message.setFromState("Offline");
  message.setToState("Slave");
  message.setPartitionName("TestDB_0");
  message.setMsgId("Some unique message id");
  message.setResourceName("TestDB");
  message.setTgtName("localhost");
  message.setStateModelDef("MasterSlave");
  message.setStateModelFactoryName(HelixConstants.DEFAULT_STATE_MODEL_FACTORY);
  MockStateModelAnnotated stateModel = new MockStateModelAnnotated();
  NotificationContext context;

  MockManager manager = new MockManager("clusterName");
  HelixDataAccessor accessor = manager.getHelixDataAccessor();

  StateModelDefinition stateModelDef =
      new StateModelDefinition(StateModelConfigGenerator.generateConfigForMasterSlave());
  Builder keyBuilder = accessor.keyBuilder();
  accessor.setProperty(keyBuilder.stateModelDef("MasterSlave"), stateModelDef);

  context = new NotificationContext(manager);

  CurrentState currentStateDelta = new CurrentState("TestDB");
  currentStateDelta.setState("TestDB_0", "OFFLINE");

  StateModelFactory<MockStateModelAnnotated> stateModelFactory =
      new StateModelFactory<MockStateModelAnnotated>() {

        @Override
        public MockStateModelAnnotated createNewStateModel(String resource, String partitionName) {
          // TODO Auto-generated method stub
          return new MockStateModelAnnotated();
        }

      };

  HelixStateTransitionHandler stHandler =
      new HelixStateTransitionHandler(stateModelFactory, stateModel, message, context,
          currentStateDelta);

  HelixTask handler = new HelixTask(message, context, stHandler, executor);
  handler.call();
  AssertJUnit.assertTrue(stateModel.stateModelInvoked);
  System.out.println("END TestCMTaskHandler.testInvocationAnnotated() at "
      + new Date(System.currentTimeMillis()));
}