TestConsumerSession.java

package com.trendmicro.mist.session;

import java.io.BufferedInputStream;
import java.io.BufferedOutputStream;
import java.io.IOException;
import java.net.Socket;
import java.net.UnknownHostException;

import javax.jms.BytesMessage;
import javax.jms.JMSException;
import javax.jms.Message;

import junit.framework.Test;
import junit.framework.TestCase;
import junit.framework.TestSuite;

import com.sun.messaging.jmq.jmsservice.JMSServiceException;
import com.trendmicro.codi.CODIException;
import com.trendmicro.codi.ZKSessionManager;
import com.trendmicro.mist.MistException;
import com.trendmicro.mist.proto.GateTalk;
import com.trendmicro.mist.proto.MistMessage;
import com.trendmicro.mist.util.OpenMQTestBroker;
import com.trendmicro.mist.util.Packet;
import com.trendmicro.mist.util.ZKTestServer;
import com.trendmicro.spn.common.util.Utils;
import com.trendmicro.tme.mfr.BrokerFarm;

public class TestConsumerSession extends TestCase {
    private ZKTestServer zkTestServer;
    private BrokerFarm brokerFarm = new BrokerFarm();

    class OnMessageRunner extends Thread {
        ConsumerSession sess;
        Message msg;

        public OnMessageRunner(ConsumerSession sess, Message msg) {
            this.sess = sess;
            this.msg = msg;
        }

        @Override
        public void run() {
            sess.onMessage(msg);
        }
    }

    private GateTalk.Session genSessionConfig(String brokerType, String host, String port, String username, String password) {
        GateTalk.Connection.Builder connBuilder = GateTalk.Connection.newBuilder();
        connBuilder.setBrokerType(brokerType);
        connBuilder.setHostName(host);
        connBuilder.setHostPort(port);
        connBuilder.setUsername(username);
        connBuilder.setPassword(password);
        return GateTalk.Session.newBuilder().setConnection(connBuilder.build()).build();
    }

    private GateTalk.Client genClientConfig(String exName) {
        GateTalk.Client.Builder builder = GateTalk.Client.newBuilder();
        builder.setChannel(GateTalk.Channel.newBuilder().setName(exName).setPersistent(false).setType(GateTalk.Channel.Type.QUEUE).build());
        builder.setAction(GateTalk.Client.Action.MOUNT);
        builder.setSessionId(0);
        builder.setType(GateTalk.Client.Type.CONSUMER);

        return builder.build();
    }

    @Override
    protected void setUp() throws Exception {
        zkTestServer = new ZKTestServer(39979);
        zkTestServer.start();
        ZKSessionManager.initialize("localhost:39979", 8000);
        super.setUp();
    }

    @Override
    protected void tearDown() throws Exception {
        ZKSessionManager.uninitialize();
        zkTestServer.stop();

        super.tearDown();
    }

    public void testOnMessage() throws MistException, UnknownHostException, IOException, ClassNotFoundException, IllegalAccessException, InstantiationException, InterruptedException, JMSException, CODIException, JMSServiceException {
        /**
         * Setup open mq
         */
        OpenMQTestBroker brk = new OpenMQTestBroker("test", 9876);
        brk.start();
        brk.registerOnZk();
        for(int i = 0; i < 10; i++) {
            if(brokerFarm.getBrokerCount() == 1)
                break;
            Utils.justSleep(500);
        }
        assertEquals(1, brokerFarm.getBrokerCount());
        assertTrue(Utils.checkSocketConnectable("localhost", 9876));

        /**
         * Session setup and attach
         */
        GateTalk.Session sessConfig = genSessionConfig("", "", "", "", "");
        ConsumerSession sess = new ConsumerSession(0, sessConfig);
        sess.attach(GateTalk.Request.Role.SOURCE);
        Socket socket = new Socket("localhost", sess.getCommPort());
        assertTrue(socket.isConnected());
        for(int i = 0; i < 10; i++) {
            if(sess.isReady())
                break;
            Utils.justSleep(500);
        }
        assertTrue(sess.isReady());

        Packet packet = new Packet();
        BufferedInputStream socketInput = new BufferedInputStream(socket.getInputStream());
        BufferedOutputStream socketOutput = new BufferedOutputStream(socket.getOutputStream());

        /**
         * Get a message and ack
         */
        brk.sendMessage(true, "foo.in", "foo".getBytes());
        BytesMessage msgin = brk.getBytesMessage(true, "foo.in");
        new OnMessageRunner(sess, msgin).start();
        packet.read(socketInput);
        MistMessage.MessageBlock mBlock = MistMessage.MessageBlock.newBuilder().mergeFrom(packet.getPayload()).build();
        assertEquals("queue:foo.in", mBlock.getId());
        assertEquals("foo", new String(mBlock.getMessage().toByteArray()));

        packet.setPayload(GateTalk.Response.newBuilder().setSuccess(true).build().toByteArray());
        packet.write(socketOutput);

        brk.stop();
    }

    public void testReceiveMessage() throws ClassNotFoundException, IllegalAccessException, InstantiationException, InterruptedException, IOException, CODIException, JMSException, MistException, JMSServiceException {
        /**
         * Setup open mq
         */
        OpenMQTestBroker brk = new OpenMQTestBroker("test", 9876);
        brk.start();
        brk.registerOnZk();
        for(int i = 0; i < 10; i++) {
            if(brokerFarm.getBrokerCount() == 1)
                break;
            Utils.justSleep(500);
        }
        assertEquals(1, brokerFarm.getBrokerCount());
        assertTrue(Utils.checkSocketConnectable("localhost", 9876));

        /**
         * Session setup and attach and mount foo.in
         */
        GateTalk.Session sessConfig = genSessionConfig("", "", "", "", "");
        ConsumerSession sess = new ConsumerSession(0, sessConfig);
        GateTalk.Client clientConfig = genClientConfig("foo.in");
        // sess.addClient(new Client(clientConfig, sessConfig));
        sess.addClient(clientConfig);
        sess.attach(GateTalk.Request.Role.SOURCE);
        Socket socket = new Socket("localhost", sess.getCommPort());
        assertTrue(socket.isConnected());
        for(int i = 0; i < 10; i++) {
            if(sess.isReady())
                break;
            Utils.justSleep(500);
        }
        assertTrue(sess.isReady());

        Packet packet = new Packet();
        BufferedInputStream socketInput = new BufferedInputStream(socket.getInputStream());
        BufferedOutputStream socketOutput = new BufferedOutputStream(socket.getOutputStream());

        /**
         * Send and receive a message from foo.in
         */
        brk.sendMessage(true, "foo.in", "foo".getBytes());
        packet.read(socketInput);
        MistMessage.MessageBlock mBlock = MistMessage.MessageBlock.newBuilder().mergeFrom(packet.getPayload()).build();
        assertEquals("queue:foo.in", mBlock.getId());
        assertEquals("foo", new String(mBlock.getMessage().toByteArray()));

        packet.setPayload(GateTalk.Response.newBuilder().setSuccess(true).build().toByteArray());
        packet.write(socketOutput);

        /**
         * Test dynamic mount bar.in
         */
        brk.sendMessage(true, "bar.in", "bar".getBytes());
        clientConfig = genClientConfig("bar.in");
        sess.addClient(clientConfig);

        packet.read(socketInput);
        mBlock = MistMessage.MessageBlock.newBuilder().mergeFrom(packet.getPayload()).build();
        assertEquals("queue:bar.in", mBlock.getId());
        assertEquals("bar", new String(mBlock.getMessage().toByteArray()));

        packet.setPayload(GateTalk.Response.newBuilder().setSuccess(true).build().toByteArray());
        packet.write(socketOutput);

        brk.stop();
    }

    public static Test suite() {
        return new TestSuite(TestConsumerSession.class);
    }
}