package com.linkedin.clustermanager;

import org.testng.AssertJUnit;
import org.testng.annotations.Test;

import com.linkedin.clustermanager.Mocks.MockManager;
import com.linkedin.clustermanager.Mocks.MockStateModel;
import com.linkedin.clustermanager.Mocks.MockStateModelAnnotated;
import com.linkedin.clustermanager.messaging.handling.CMStateTransitionHandler;
import com.linkedin.clustermanager.messaging.handling.CMTaskHandler;
import com.linkedin.clustermanager.model.Message;
import com.linkedin.clustermanager.model.Message.MessageType;

public class TestCMTaskHandler
{
  @Test
  public void testInvocation() throws Exception
  {
    System.out.println("TestCMTaskHandler.testInvocation()");
    Message message = new Message(MessageType.STATE_TRANSITION);
    message.setSrcName("cm-instance-0");
    message.setTgtSessionId("1234");
    message.setFromState("Offline");
    message.setToState("Slave");
    message.setStateUnitKey("Teststateunitkey");
    message.setId("Some unique id");
    message.setMsgId("Some unique message id");
    message.setStateUnitGroup("TeststateunitGroup");
    MockStateModel stateModel = new MockStateModel();
    NotificationContext context;

    context = new NotificationContext(new MockManager());
    CMTaskHandler handler;
    handler = new CMTaskHandler(message, context, null, null);
    handler.call();
    AssertJUnit.assertTrue(stateModel.stateModelInvoked);
  }

  @Test
  public void testInvocationAnnotated() throws Exception
  {
    System.out.println("TestCMTaskHandler.testInvocationAnnotated()");
    Message message = new Message(MessageType.STATE_TRANSITION);
    message.setSrcName("cm-instance-0");
    message.setTgtSessionId("1234");
    message.setFromState("Offline");
    message.setToState("Slave");
    message.setStateUnitKey("Teststateunitkey");
    message.setId("Some unique id");
    message.setMsgId("Some unique message id");
    MockStateModelAnnotated stateModel = new MockStateModelAnnotated();
    NotificationContext context;

    context = new NotificationContext(new MockManager());
    CMTaskHandler handler;
    handler = new CMTaskHandler(message, context, null, null);
    handler.call();
    AssertJUnit.assertTrue(stateModel.stateModelInvoked);
  }

}

@@ -1171,32 +1171,115 @@
 Context context;
+%0A    CMStateTransitionHandler stHandler = new CMStateTransitionHandler(stateModel);
 %0A%0A    context = 
@@ -1391,36 +1391,41 @@
 ssage, context, 
-null
+stHandler
 , null);%0A    han
@@ -2170,32 +2170,116 @@
 andler handler;%0A
+    CMStateTransitionHandler stHandler = new CMStateTransitionHandler(stateModel);%0A%0A
     handler = ne
@@ -2312,20 +2312,25 @@
 ontext, 
-null
+stHandler
 , null);

Test change since we have added state transition handler