package com.linkedin.clustermanager;

import java.util.HashSet;
import java.util.Map;
import java.util.UUID;

import org.testng.Assert;
import org.testng.annotations.Test;

import com.linkedin.clustermanager.messaging.handling.MessageHandler;
import com.linkedin.clustermanager.messaging.handling.MessageHandlerFactory;
import com.linkedin.clustermanager.model.Message;
import com.linkedin.clustermanager.model.Message.MessageType;
import com.linkedin.clustermanager.tools.ClusterStateVerifier;

public class TestMessagingService extends ZkStandAloneCMHandler
{
  public static class TestMessagingHandlerFactory implements MessageHandlerFactory
  {
    public static HashSet<String> _processedMsgIds = new HashSet<String>();
    @Override
    public MessageHandler createHandler(Message message,
        NotificationContext context)
    {
      return new TestMessagingHandler();
    }

    @Override
    public String getMessageType()
    {
      return "TestExtensibility";
    }

    @Override
    public void reset()
    {
      // TODO Auto-generated method stub
      
    }
    
    public static class TestMessagingHandler implements MessageHandler
    {
      @Override
      public void handleMessage(Message message, NotificationContext context,
          Map<String, String> resultMap) throws InterruptedException
      {
        // TODO Auto-generated method stub
        System.out.println("TestMessagingHandler " + message.getMsgId());
        _processedMsgIds.add(message.getRecord().getSimpleField("TestMessagingPara"));
        resultMap.put("ReplyMessage", "TestReplyMessage");
      }
    }
  }
  
  
  @Test
  public void TestMessageSimpleSend() throws Exception
  {
    String hostSrc = "localhost_"+START_PORT;
    String hostDest = "localhost_"+(START_PORT + 1);
    
    TestMessagingHandlerFactory factory = new TestMessagingHandlerFactory();
    _managerMap.get(hostDest).getMessagingService().registerMessageHandlerFactory(factory.getMessageType(), factory);
    
    String msgId = new UUID(123,456).toString(); 
    Message msg = new Message(factory.getMessageType());
    msg.setMsgId(msgId);
    msg.setSrcName(hostSrc);
    msg.setTgtSessionId("*");
    msg.setMsgState("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);
    
    _managerMap.get(hostSrc).getMessagingService().send(cr, msg);
    
    Thread.currentThread().sleep(2000);
    //Thread.currentThread().join();
    Assert.assertTrue(TestMessagingHandlerFactory._processedMsgIds.contains(para));
    
  }
  public static class TestAsyncCallback extends AsyncCallback
  {
    static HashSet<String> _replyedMessageContents = new HashSet<String>();
    @Override
    public void onTimeOut()
    {
      // TODO Auto-generated method stub
      
    }

    @Override
    public void onReplyMessage(Message message)
    {
      // TODO Auto-generated method stub
      System.out.println("OnreplyMessage: "+ message.getRecord().getMapField(Message.Attributes.MESSAGE_RESULT.toString()).get("ReplyMessage"));
      _replyedMessageContents.add(message.getRecord().getMapField(Message.Attributes.MESSAGE_RESULT.toString()).get("ReplyMessage"));
    }
    
  }
  @Test
  public void TestMessageSimpleSendReceiveAsync() throws Exception
  {
    String hostSrc = "localhost_"+START_PORT;
    String hostDest = "localhost_"+(START_PORT + 1);
    
    TestMessagingHandlerFactory factory = new TestMessagingHandlerFactory();
    _managerMap.get(hostDest).getMessagingService().registerMessageHandlerFactory(factory.getMessageType(), factory);
    
    
    _managerMap.get(hostSrc).getMessagingService().registerMessageHandlerFactory(factory.getMessageType(), factory);
    
    
    String msgId = new UUID(123,456).toString(); 
    Message msg = new Message(factory.getMessageType());
    msg.setMsgId(msgId);
    msg.setSrcName(hostSrc);
    
    msg.setTgtSessionId("*");
    msg.setMsgState("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);
    
    TestAsyncCallback callback = new TestAsyncCallback();
    
    _managerMap.get(hostSrc).getMessagingService().send(cr, msg, callback);
    
    Thread.currentThread().sleep(2000);
    //Thread.currentThread().join();
    Assert.assertTrue(TestAsyncCallback._replyedMessageContents.contains("TestReplyMessage"));
    
  }
  
  @Test
  public void TestBlockingSendReceive() throws Exception
  {
    String hostSrc = "localhost_"+START_PORT;
    String hostDest = "localhost_"+(START_PORT + 1);
    
    TestMessagingHandlerFactory factory = new TestMessagingHandlerFactory();
    _managerMap.get(hostDest).getMessagingService().registerMessageHandlerFactory(factory.getMessageType(), factory);
    
    String msgId = new UUID(123,456).toString(); 
    Message msg = new Message(factory.getMessageType());
    msg.setMsgId(msgId);
    msg.setSrcName(hostSrc);
    
    msg.setTgtSessionId("*");
    msg.setMsgState("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 result = _managerMap.get(hostSrc).getMessagingService().sendReceive(cr, msg, 2000);
    
    Assert.assertTrue(result.getMessageReplied().get(0).getRecord().getMapField(Message.Attributes.MESSAGE_RESULT.toString()).get("ReplyMessage").equals("TestReplyMessage"));
    
  }
}

@@ -1366,24 +1366,68 @@
 method stub%0A
+        Thread.currentThread().sleep(2000);%0A
         Syst
@@ -1484,24 +1484,24 @@
 etMsgId());%0A
-
         _pro
@@ -2623,33 +2623,33 @@
 Thread().sleep(2
-0
+5
 00);%0A    //Threa
@@ -2901,32 +2901,68 @@
 hSet%3CString%3E();%0A
+    public boolean timeout = false;%0A
     @Override%0A  
@@ -3036,24 +3036,39 @@
  stub%0A      
+timeout = true;
 %0A    %7D%0A%0A    
@@ -4652,33 +4652,33 @@
 tThread().sleep(
-2
+3
 000);%0A    //Thre
@@ -4800,24 +4800,262 @@
 ge%22));%0A    %0A
+    callback.setTimeout(500);%0A    _managerMap.get(hostSrc).getMessagingService().send(cr, msg, callback);%0A    %0A    Thread.currentThread().sleep(3000);%0A    //Thread.currentThread().join();%0A    Assert.assertTrue(callback.isTimeOut());%0A    %0A
   %7D%0A  %0A  @Te
@@ -6024,17 +6024,17 @@
 r, msg, 
-2
+3
 000);%0A  
@@ -6209,21 +6209,155 @@
 ssage%22));%0A  
-  
+%0A    result = _managerMap.get(hostSrc).getMessagingService().sendReceive(cr, msg, 1000);%0A    %0A    Assert.assertTrue(result.isTimeOut());
 %0A  %7D%0A%7D%0A

Add test case for message API timeout