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

@Test(dependsOnMethods = "testIsInstanceStoppable")
public void testGetAllMessages() throws IOException {
  System.out.println("Start test :" + TestHelper.getTestMethodName());
  String testInstance = CLUSTER_NAME + "localhost_12926"; //Non-live instance

  String messageId = "msg1";
  Message message = new Message(Message.MessageType.STATE_TRANSITION, messageId);
  message.setStateModelDef("MasterSlave");
  message.setFromState("OFFLINE");
  message.setToState("SLAVE");
  message.setResourceName("testResourceName");
  message.setPartitionName("testResourceName_1");
  message.setTgtName("localhost_3");
  message.setTgtSessionId("session_3");
  HelixDataAccessor helixDataAccessor = new ZKHelixDataAccessor(CLUSTER_NAME, _baseAccessor);
  helixDataAccessor.setProperty(helixDataAccessor.keyBuilder().message(testInstance, messageId), message);

  String body = new JerseyUriRequestBuilder("clusters/{}/instances/{}/messages")
      .isBodyReturnExpected(true).format(CLUSTER_NAME, testInstance).get(this);
  JsonNode node = OBJECT_MAPPER.readTree(body);
  int newMessageCount =
      node.get(PerInstanceAccessor.PerInstanceProperties.total_message_count.name()).getIntValue();

  Assert.assertEquals(newMessageCount, 1);
  System.out.println("End test :" + TestHelper.getTestMethodName());
}
 

private Message generateMessage(String from, String to) {
  String uuid = UUID.randomUUID().toString();
  Message message = new Message(Message.MessageType.STATE_TRANSITION, uuid);
  message.setSrcName("ADMIN");
  message.setTgtName(_participant.getInstanceName());
  message.setMsgState(Message.MessageState.NEW);
  message.setPartitionName("P");
  message.setResourceName(WorkflowGenerator.DEFAULT_TGT_DB);
  message.setFromState(from);
  message.setToState(to);
  message.setTgtSessionId(_participant.getSessionId());
  message.setSrcSessionId(_manager.getSessionId());
  message.setStateModelDef("OnlineOffline");
  message.setStateModelFactoryName("DEFAULT");
  return message;
}
 

private Message newMessage(String resourceName, String partitionName, String instanceName,
    String fromState, String toState) {
  String uuid = UUID.randomUUID().toString();
  Message message = new Message(MessageType.STATE_TRANSITION, uuid);
  message.setSrcName("controller");
  message.setTgtName(instanceName);
  message.setMsgState(MessageState.NEW);
  message.setResourceName(resourceName);
  message.setPartitionName(partitionName);
  message.setFromState(fromState);
  message.setToState(toState);
  message.setTgtSessionId("sessionId");
  message.setSrcSessionId("sessionId");
  message.setStateModelDef("MasterSlave");
  message.setStateModelFactoryName("DEFAULT");
  message.setBucketSize(0);
  return message;
}
 

private Message createStateTransitionCancellationMessage(HelixManager manager, Resource resource,
    String partitionName, String instanceName, String sessionId, String stateModelDefName,
    String fromState, String toState, String nextState, Message cancellationMessage,
    boolean isCancellationEnabled, String currentState) {

  if (isCancellationEnabled && cancellationMessage == null) {
    LogUtil.logInfo(logger, _eventId,
        "Send cancellation message of the state transition for " + resource.getResourceName()
            + "." + partitionName + " on " + instanceName + ", currentState: " + currentState
            + ", nextState: " + (nextState == null ? "N/A" : nextState));

    String uuid = UUID.randomUUID().toString();
    Message message = new Message(MessageType.STATE_TRANSITION_CANCELLATION, uuid);
    message.setSrcName(manager.getInstanceName());
    message.setTgtName(instanceName);
    message.setMsgState(MessageState.NEW);
    message.setPartitionName(partitionName);
    message.setResourceName(resource.getResourceName());
    message.setFromState(fromState);
    message.setToState(toState);
    message.setTgtSessionId(sessionId);
    message.setSrcSessionId(manager.getSessionId());
    message.setStateModelDef(stateModelDefName);
    message.setStateModelFactoryName(resource.getStateModelFactoryname());
    message.setBucketSize(resource.getBucketSize());
    return message;
  }

  return null;
}
 

@Test(dependsOnMethods = "testGetAllMessages")
public void testGetMessagesByStateModelDef() throws IOException {
  System.out.println("Start test :" + TestHelper.getTestMethodName());

  String testInstance = CLUSTER_NAME + "localhost_12926"; //Non-live instance
  String messageId = "msg1";
  Message message = new Message(Message.MessageType.STATE_TRANSITION, messageId);
  message.setStateModelDef("MasterSlave");
  message.setFromState("OFFLINE");
  message.setToState("SLAVE");
  message.setResourceName("testResourceName");
  message.setPartitionName("testResourceName_1");
  message.setTgtName("localhost_3");
  message.setTgtSessionId("session_3");
  HelixDataAccessor helixDataAccessor = new ZKHelixDataAccessor(CLUSTER_NAME, _baseAccessor);
  helixDataAccessor.setProperty(helixDataAccessor.keyBuilder().message(testInstance, messageId),
      message);

  String body =
      new JerseyUriRequestBuilder("clusters/{}/instances/{}/messages?stateModelDef=MasterSlave")
          .isBodyReturnExpected(true).format(CLUSTER_NAME, testInstance).get(this);
  JsonNode node = OBJECT_MAPPER.readTree(body);
  int newMessageCount =
      node.get(PerInstanceAccessor.PerInstanceProperties.total_message_count.name()).getIntValue();

  Assert.assertEquals(newMessageCount, 1);

  body =
      new JerseyUriRequestBuilder("clusters/{}/instances/{}/messages?stateModelDef=LeaderStandBy")
          .isBodyReturnExpected(true).format(CLUSTER_NAME, testInstance).get(this);
  node = OBJECT_MAPPER.readTree(body);
  newMessageCount =
      node.get(PerInstanceAccessor.PerInstanceProperties.total_message_count.name()).getIntValue();

  Assert.assertEquals(newMessageCount, 0);
  System.out.println("End test :" + TestHelper.getTestMethodName());
}
 

@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()));
}
 

public static Message createMessage(String msgId, String fromState, String toState,
    String tgtName, String resourceName, String partitionName) {
  Message msg = new Message(MessageType.STATE_TRANSITION, msgId);
  msg.setFromState(fromState);
  msg.setToState(toState);
  msg.setTgtName(tgtName);
  msg.setResourceName(resourceName);
  msg.setPartitionName(partitionName);
  msg.setStateModelDef("MasterSlave");

  return msg;
}
 

private List<Message> generateMessagesForParticipant(Criteria recipientCriteria, Message message,
    HelixDataAccessor targetDataAccessor) {
  List<Message> messages = new ArrayList<Message>();
  List<Map<String, String>> matchedList =
      _evaluator.evaluateCriteria(recipientCriteria, targetDataAccessor);

  if (!matchedList.isEmpty()) {
    Map<String, String> sessionIdMap = new HashMap<String, String>();
    if (recipientCriteria.isSessionSpecific()) {
      Builder keyBuilder = targetDataAccessor.keyBuilder();
      // For backward compatibility, allow partial read for the live instances.
      // Note that this may cause the pending message to be sent with null target session Id.
      List<LiveInstance> liveInstances =
          targetDataAccessor.getChildValues(keyBuilder.liveInstances(), false);

      for (LiveInstance liveInstance : liveInstances) {
        sessionIdMap.put(liveInstance.getInstanceName(), liveInstance.getEphemeralOwner());
      }
    }
    for (Map<String, String> map : matchedList) {
      String id = UUID.randomUUID().toString();
      Message newMessage = new Message(message.getRecord(), id);
      String srcInstanceName = _manager.getInstanceName();
      String tgtInstanceName = map.get("instanceName");
      // Don't send message to self
      if (recipientCriteria.isSelfExcluded() && srcInstanceName.equalsIgnoreCase(tgtInstanceName)) {
        continue;
      }
      newMessage.setSrcName(srcInstanceName);
      newMessage.setTgtName(tgtInstanceName);
      newMessage.setResourceName(map.get("resourceName"));
      newMessage.setPartitionName(map.get("partitionName"));
      if (recipientCriteria.isSessionSpecific()) {
        newMessage.setTgtSessionId(sessionIdMap.get(tgtInstanceName));
      }
      messages.add(newMessage);
    }
  }
  return messages;
}
 

private Message createStateTransitionMessage(HelixManager manager, Resource resource,
    String partitionName, String instanceName, String currentState, String nextState,
    String sessionId, String stateModelDefName) {
  String uuid = UUID.randomUUID().toString();
  Message message = new Message(MessageType.STATE_TRANSITION, uuid);
  message.setSrcName(manager.getInstanceName());
  message.setTgtName(instanceName);
  message.setMsgState(MessageState.NEW);
  message.setPartitionName(partitionName);
  message.setResourceName(resource.getResourceName());
  message.setFromState(currentState);
  message.setToState(nextState);
  message.setTgtSessionId(sessionId);
  message.setSrcSessionId(manager.getSessionId());
  message.setStateModelDef(stateModelDefName);
  message.setStateModelFactoryName(resource.getStateModelFactoryname());
  message.setBucketSize(resource.getBucketSize());

  if (resource.getResourceGroupName() != null) {
    message.setResourceGroupName(resource.getResourceGroupName());
  }
  if (resource.getResourceTag() != null) {
    message.setResourceTag(resource.getResourceTag());
  }

  return message;
}
 

@Test()
public void TestParticipantErrorMessageSend() {
  String participant1 = "localhost_" + START_PORT;
  String participant2 = "localhost_" + (START_PORT + 1);

  Message errorMessage1 =
      new Message(MessageType.PARTICIPANT_ERROR_REPORT, UUID.randomUUID().toString());
  errorMessage1.setTgtSessionId("*");
  errorMessage1.getRecord().setSimpleField(
      DefaultParticipantErrorMessageHandlerFactory.ACTIONKEY,
      ActionOnError.DISABLE_INSTANCE.toString());
  Criteria recipientCriteria = new Criteria();
  recipientCriteria.setRecipientInstanceType(InstanceType.CONTROLLER);
  recipientCriteria.setSessionSpecific(false);
  _participants[0].getMessagingService().send(recipientCriteria,
      errorMessage1);

  Message errorMessage2 =
      new Message(MessageType.PARTICIPANT_ERROR_REPORT, UUID.randomUUID().toString());
  errorMessage2.setTgtSessionId("*");
  errorMessage2.setResourceName("TestDB");
  errorMessage2.setPartitionName("TestDB_14");
  errorMessage2.getRecord().setSimpleField(
      DefaultParticipantErrorMessageHandlerFactory.ACTIONKEY,
      ActionOnError.DISABLE_PARTITION.toString());
  Criteria recipientCriteria2 = new Criteria();
  recipientCriteria2.setRecipientInstanceType(InstanceType.CONTROLLER);
  recipientCriteria2.setSessionSpecific(false);
  _participants[1].getMessagingService().send(recipientCriteria2,
      errorMessage2);

  try {
    Thread.sleep(1500);
  } catch (InterruptedException e) {
    LOG.error("Interrupted sleep", e);
  }

  boolean result =
      ClusterStateVerifier.verifyByZkCallback(new BestPossAndExtViewZkVerifier(ZK_ADDR,
          CLUSTER_NAME));
  Assert.assertTrue(result);
  Builder kb = _participants[1].getHelixDataAccessor().keyBuilder();
  ExternalView externalView =
      _participants[1].getHelixDataAccessor().getProperty(
          kb.externalView("TestDB"));

  for (String partitionName : externalView.getRecord().getMapFields().keySet()) {
    for (String hostName : externalView.getRecord().getMapField(partitionName).keySet()) {
      if (hostName.equals(participant1)) {
        Assert.assertTrue(externalView.getRecord().getMapField(partitionName).get(hostName)
            .equalsIgnoreCase("OFFLINE"));
      }
    }
  }
  Assert.assertTrue(externalView.getRecord().getMapField("TestDB_14").get(participant2)
      .equalsIgnoreCase("OFFLINE"));
}
 

@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 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 testInvocation() throws Exception {

  HelixManager testHelixManager =
      HelixManagerFactory.getZKHelixManager(clusterName, "localhost_8900",
          InstanceType.PARTICIPANT, ZK_ADDR);
  testHelixManager.connect();

  TestZKCallback test = new TestZKCallback();

  TestZKCallback.TestCallbackListener testListener = test.new TestCallbackListener();

  testHelixManager.addMessageListener(testListener, "localhost_8900");
  testHelixManager.addCurrentStateChangeListener(testListener, "localhost_8900",
      testHelixManager.getSessionId());
  testHelixManager.addCustomizedStateRootChangeListener(testListener, "localhost_8900");
  testHelixManager.addConfigChangeListener(testListener);
  testHelixManager.addIdealStateChangeListener(testListener);
  testHelixManager.addExternalViewChangeListener(testListener);
  testHelixManager.addLiveInstanceChangeListener(testListener);
  // Initial add listener should trigger the first execution of the
  // listener callbacks
  AssertJUnit.assertTrue(testListener.configChangeReceived
      & testListener.currentStateChangeReceived & testListener.externalViewChangeReceived
      & testListener.idealStateChangeReceived & testListener.liveInstanceChangeReceived
      & testListener.messageChangeReceived);

  testListener.Reset();
  HelixDataAccessor accessor = testHelixManager.getHelixDataAccessor();
  Builder keyBuilder = accessor.keyBuilder();

  ExternalView extView = new ExternalView("db-12345");
  accessor.setProperty(keyBuilder.externalView("db-12345"), extView);
  Thread.sleep(100);
  AssertJUnit.assertTrue(testListener.externalViewChangeReceived);
  testListener.Reset();

  CurrentState curState = new CurrentState("db-12345");
  curState.setSessionId("sessionId");
  curState.setStateModelDefRef("StateModelDef");
  accessor.setProperty(keyBuilder.currentState("localhost_8900", testHelixManager.getSessionId(),
      curState.getId()), curState);
  Thread.sleep(100);
  AssertJUnit.assertTrue(testListener.currentStateChangeReceived);
  testListener.Reset();

  IdealState idealState = new IdealState("db-1234");
  idealState.setNumPartitions(400);
  idealState.setReplicas(Integer.toString(2));
  idealState.setStateModelDefRef("StateModeldef");
  accessor.setProperty(keyBuilder.idealStates("db-1234"), idealState);
  Thread.sleep(100);
  AssertJUnit.assertTrue(testListener.idealStateChangeReceived);
  testListener.Reset();

  // dummyRecord = new ZNRecord("db-12345");
  // dataAccessor.setProperty(PropertyType.IDEALSTATES, idealState, "db-12345"
  // );
  // Thread.sleep(100);
  // AssertJUnit.assertTrue(testListener.idealStateChangeReceived);
  // testListener.Reset();

  // dummyRecord = new ZNRecord("localhost:8900");
  // List<ZNRecord> recList = new ArrayList<ZNRecord>();
  // recList.add(dummyRecord);

  testListener.Reset();
  Message message = new Message(MessageType.STATE_TRANSITION, UUID.randomUUID().toString());
  message.setTgtSessionId("*");
  message.setResourceName("testResource");
  message.setPartitionName("testPartitionKey");
  message.setStateModelDef("MasterSlave");
  message.setToState("toState");
  message.setFromState("fromState");
  message.setTgtName("testTarget");
  message.setStateModelFactoryName(HelixConstants.DEFAULT_STATE_MODEL_FACTORY);

  accessor.setProperty(keyBuilder.message("localhost_8900", message.getId()), message);
  Thread.sleep(500);
  AssertJUnit.assertTrue(testListener.messageChangeReceived);

  // dummyRecord = new ZNRecord("localhost_9801");
  LiveInstance liveInstance = new LiveInstance("localhost_9801");
  liveInstance.setSessionId(UUID.randomUUID().toString());
  liveInstance.setHelixVersion(UUID.randomUUID().toString());
  accessor.setProperty(keyBuilder.liveInstance("localhost_9801"), liveInstance);
  Thread.sleep(500);
  AssertJUnit.assertTrue(testListener.liveInstanceChangeReceived);
  testListener.Reset();

  // dataAccessor.setNodeConfigs(recList); Thread.sleep(100);
  // AssertJUnit.assertTrue(testListener.configChangeReceived);
  // testListener.Reset();

  accessor.removeProperty(keyBuilder.liveInstance("localhost_8900"));
  accessor.removeProperty(keyBuilder.liveInstance("localhost_9801"));
}
 

@Test()
public void testDuplicatedMessage() throws InterruptedException {
  System.out.println("START TestHelixTaskExecutor.testDuplicatedMessage()");
  HelixTaskExecutor executor = new HelixTaskExecutor();
  HelixManager manager = new MockClusterManager();
  HelixDataAccessor dataAccessor = manager.getHelixDataAccessor();
  PropertyKey.Builder keyBuilder = dataAccessor.keyBuilder();

  TestStateTransitionHandlerFactory stateTransitionFactory =
      new TestStateTransitionHandlerFactory(Message.MessageType.STATE_TRANSITION.name(), 1000);
  executor.registerMessageHandlerFactory(Message.MessageType.STATE_TRANSITION.name(),
      stateTransitionFactory);

  NotificationContext changeContext = new NotificationContext(manager);
  List<Message> msgList = new ArrayList<Message>();

  int nMsgs = 3;
  String instanceName = manager.getInstanceName();
  for (int i = 0; i < nMsgs; i++) {
    Message msg =
        new Message(Message.MessageType.STATE_TRANSITION.name(), UUID.randomUUID().toString());
    msg.setTgtSessionId(manager.getSessionId());
    msg.setCreateTimeStamp((long) i);
    msg.setTgtName("Localhost_1123");
    msg.setSrcName("127.101.1.23_2234");
    msg.setPartitionName("Partition");
    msg.setResourceName("Resource");
    msg.setStateModelDef("DummyMasterSlave");
    msg.setFromState("SLAVE");
    msg.setToState("MASTER");
    dataAccessor.setProperty(msg.getKey(keyBuilder, instanceName), msg);
    msgList.add(msg);
  }

  AssertJUnit
      .assertEquals(dataAccessor.getChildValues(keyBuilder.messages(instanceName), true).size(),
          nMsgs);

  changeContext.setChangeType(HelixConstants.ChangeType.MESSAGE);
  executor.onMessage(instanceName, msgList, changeContext);

  Thread.sleep(200);

  // only 1 message is left over - state transition takes 1sec
  Assert.assertEquals(dataAccessor.getChildValues(keyBuilder.messages(instanceName), true).size(),
      1);

  // While a state transition message is going on, another state transition message for same
  // resource / partition comes in, it should be discarded by message handler

  // Mock accessor is modifying message state in memory so we set it back to NEW
  msgList.get(2).setMsgState(MessageState.NEW);
  dataAccessor.setProperty(msgList.get(2).getKey(keyBuilder, instanceName), msgList.get(2));
  executor.onMessage(instanceName, Arrays.asList(msgList.get(2)), changeContext);
  Thread.sleep(200);
  Assert.assertEquals(dataAccessor.getChildValues(keyBuilder.messages(instanceName), true).size(),
      1);

  Thread.sleep(1000);
  Assert.assertEquals(dataAccessor.getChildValues(keyBuilder.messages(instanceName), true).size(),
      0);
  System.out.println("END TestHelixTaskExecutor.testDuplicatedMessage()");
}
 

@Test
public void testStateTransitionCancellationMsg() throws InterruptedException {
  HelixTaskExecutor executor = new HelixTaskExecutor();
  HelixManager manager = new MockClusterManager();

  TestStateTransitionHandlerFactory stateTransitionFactory = new TestStateTransitionHandlerFactory(Message.MessageType.STATE_TRANSITION.name());
  TestStateTransitionHandlerFactory cancelFactory = new TestStateTransitionHandlerFactory(Message.MessageType.STATE_TRANSITION_CANCELLATION
      .name());
  executor.registerMessageHandlerFactory(Message.MessageType.STATE_TRANSITION.name(), stateTransitionFactory);
  executor.registerMessageHandlerFactory(Message.MessageType.STATE_TRANSITION_CANCELLATION.name(), cancelFactory);


  NotificationContext changeContext = new NotificationContext(manager);

  List<Message> msgList = new ArrayList<Message>();
  Message msg1 = new Message(Message.MessageType.STATE_TRANSITION, UUID.randomUUID().toString());
  msg1.setTgtSessionId("*");
  msg1.setPartitionName("P1");
  msg1.setResourceName("R1");
  msg1.setTgtName("Localhost_1123");
  msg1.setSrcName("127.101.1.23_2234");
  msg1.setFromState("SLAVE");
  msg1.setToState("MASTER");
  msgList.add(msg1);

  Message msg2 = new Message(Message.MessageType.STATE_TRANSITION_CANCELLATION, UUID.randomUUID().toString());
  msg2.setTgtSessionId("*");
  msg2.setPartitionName("P1");
  msg2.setResourceName("R1");
  msg2.setTgtName("Localhost_1123");
  msg2.setSrcName("127.101.1.23_2234");
  msg2.setFromState("SLAVE");
  msg2.setToState("MASTER");
  msgList.add(msg2);

  changeContext.setChangeType(HelixConstants.ChangeType.MESSAGE);
  executor.onMessage("someInstance", msgList, changeContext);

  Thread.sleep(3000);
  AssertJUnit.assertEquals(cancelFactory._processedMsgIds.size(), 0);
  AssertJUnit.assertEquals(stateTransitionFactory._processedMsgIds.size(), 0);
}
 

@Test
public void testStateMismatch() throws InterruptedException {
  // String controllerName = CONTROLLER_PREFIX + "_0";

  HelixManager manager = _controller; // _startCMResultMap.get(controllerName)._manager;
  HelixDataAccessor accessor = manager.getHelixDataAccessor();
  Builder kb = accessor.keyBuilder();
  ExternalView ev = accessor.getProperty(kb.externalView(TEST_DB));
  Map<String, LiveInstance> liveinstanceMap =
      accessor.getChildValuesMap(accessor.keyBuilder().liveInstances(), true);

  for (String instanceName : liveinstanceMap.keySet()) {
    String sessionid = liveinstanceMap.get(instanceName).getEphemeralOwner();
    for (String partition : ev.getPartitionSet()) {
      if (ev.getStateMap(partition).containsKey(instanceName)) {
        String uuid = UUID.randomUUID().toString();
        Message message = new Message(MessageType.STATE_TRANSITION, uuid);
        boolean rand = new Random().nextInt(10) > 5;
        if (ev.getStateMap(partition).get(instanceName).equals("MASTER")) {
          message.setSrcName(manager.getInstanceName());
          message.setTgtName(instanceName);
          message.setMsgState(MessageState.NEW);
          message.setPartitionName(partition);
          message.setResourceName(TEST_DB);
          message.setFromState(rand ? "SLAVE" : "OFFLINE");
          message.setToState(rand ? "MASTER" : "SLAVE");
          message.setTgtSessionId(sessionid);
          message.setSrcSessionId(manager.getSessionId());
          message.setStateModelDef("MasterSlave");
          message.setStateModelFactoryName("DEFAULT");
        } else if (ev.getStateMap(partition).get(instanceName).equals("SLAVE")) {
          message.setSrcName(manager.getInstanceName());
          message.setTgtName(instanceName);
          message.setMsgState(MessageState.NEW);
          message.setPartitionName(partition);
          message.setResourceName(TEST_DB);
          message.setFromState(rand ? "MASTER" : "OFFLINE");
          message.setToState(rand ? "SLAVE" : "SLAVE");
          message.setTgtSessionId(sessionid);
          message.setSrcSessionId(manager.getSessionId());
          message.setStateModelDef("MasterSlave");
          message.setStateModelFactoryName("DEFAULT");
        }
        accessor.setProperty(accessor.keyBuilder().message(instanceName, message.getMsgId()),
            message);
      }
    }
  }
  Thread.sleep(3000);
  ExternalView ev2 = accessor.getProperty(kb.externalView(TEST_DB));
  Assert.assertTrue(ev.equals(ev2));
}
 

@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()));
}