Package

Source Code of AbstractStormWithKafkaTest

/*
* Author: cbedford
* Date: 11/1/13
* Time: 5:00 PM
*/


import backtype.storm.Config;
import backtype.storm.LocalCluster;
import backtype.storm.generated.StormTopology;
import com.google.common.io.Files;
import kafka.admin.CreateTopicCommand;
import kafka.server.KafkaConfig;
import kafka.server.KafkaServer;
import kafka.utils.MockTime;
import org.testng.annotations.AfterClass;
import org.testng.annotations.BeforeClass;

import java.io.File;
import java.util.Properties;
import java.util.Timer;
import java.util.concurrent.CountDownLatch;


/**
* Simplifies testing of Storm components that consume or produce data items from or to Kafka.
* Operates via  a 'template method' series of steps, wherein the BeforeClass method sets up a
* Storm Local cluster, then waits for the zookeeper instance started by that cluster to 'boot up',
* then starts an-process Kafka server using that zookeeper, and then creates a topic whose
* name is derived from the name of the base class test.
* <p/>
* Subclasses only need to implement the abstract createTopology() method (and perhaps
* override 'verifyResults())' which is currently kind of hard coded to our first two subclasses of
* this base class.
*/
public abstract class AbstractStormWithKafkaTest {
    public static String[] sentences = new String[]{
            "one dog9 - saw the fox over the moon",
            "two cats9 - saw the fox over the moon",
            "four bears9 - saw the fox over the moon",
            "five goats9 - saw the fox over the moon",
            "SHUTDOWN",
    };
    protected final String BROKER_CONNECT_STRING = "localhost:9092";    // kakfa broker server/port info
    private final String topicName = this.getClass().getSimpleName() + "_topic_" + getRandomInteger(1000);
    protected final String topologyName = this.getClass().getSimpleName() + "-topology" + getRandomInteger(1000);

    protected LocalCluster cluster = null;

    private final File kafkaWorkingDir = Files.createTempDir();
    private final CountDownLatch kafkaTopicCreatedLatch = new CountDownLatch(1);
    private KafkaServer kafkaServer = null;
    private Timer timer;
    private Thread kafkaServerThread = null;

    @BeforeClass(alwaysRun = true)
    protected void setUp() {
        timer = ServerAndThreadCoordinationUtils.setMaxTimeToRunTimer(getMaxAllowedToRunMillisecs());
        ServerAndThreadCoordinationUtils.removePauseSentinelFile();
        cluster = new LocalCluster();
        ServerAndThreadCoordinationUtils.waitForServerUp("localhost", 2000, 5 * KafkaOutputBoltTest.SECOND);   // Wait for zookeeper to come up

        /*
         *  Below we start up kafka and create topic in a separate thread. If we don't do this then we
         *  get very bizarre behavior, such as tuples never being emitted from our spouts and bolts
         *  as expected. Haven't figure out why this is needed... But doing it 'cause that's what makes
         *  things work.
         */
        kafkaServerThread = new Thread(
                new Runnable() {
                    @Override
                    public void run() {
                        startKafkaServer();
                        createTopic(getTopicName());
                        if (getSecondTopicName() != null) {
                            createTopic(getSecondTopicName());
                        }
                        ServerAndThreadCoordinationUtils.countDown(kafkaTopicCreatedLatch);
                    }
                },
                "kafkaServerThread"
        );
        kafkaServerThread.start();
        ServerAndThreadCoordinationUtils.await(kafkaTopicCreatedLatch);
    }


    public String getSecondTopicName() {
        return null;
    }


    abstract protected int getMaxAllowedToRunMillisecs();

    @AfterClass(alwaysRun = true)
    protected void tearDown() {
        try {
            kafkaServerThread.join();
        } catch (InterruptedException e) {
            e.printStackTrace();
        }

        cluster.shutdown();
        kafkaServer.shutdown();
        timer.cancel();
    }

    protected void createTopic(String topicName) {
        String[] arguments = new String[6];
        arguments[0] = "--zookeeper";
        arguments[1] = "localhost:2000";
        arguments[2] = "--partition";
        arguments[3] = "1";
        arguments[4] = "--topic";
        arguments[5] = topicName;

        CreateTopicCommand.main(arguments);
    }

    protected void startKafkaServer() {
        Properties props = createProperties(kafkaWorkingDir.getAbsolutePath(), 9092, 1);
        KafkaConfig kafkaConfig = new KafkaConfig(props);

        kafkaServer = new KafkaServer(kafkaConfig, new MockTime());
        kafkaServer.startup();
    }

    protected String getZkConnect() {   // Uses zookeeper created by LocalCluster

        return "localhost:2000";
    }

    protected int getRandomInteger(int max) {
        return (int) Math.floor((Math.random() * max));
    }

    private Properties createProperties(String logDir, int port, int brokerId) {
        Properties properties = new Properties();
        properties.put("port", port + "");
        properties.put("broker.id", brokerId + "");
        properties.put("log.dir", logDir);
        properties.put("zookeeper.connect", getZkConnect());
        return properties;
    }


    protected abstract StormTopology createTopology();


    /**
     * @return a Config object with time outs set very high so that the storm to zookeeper
     *         session will be kept alive, even as we are rooting around in a debugger.
     */
    public static Config getDebugConfigForStormTopology() {
        Config config = new Config();
        config.setDebug(true);
        config.put(Config.STORM_ZOOKEEPER_CONNECTION_TIMEOUT, 900 * 1000);
        config.put(Config.STORM_ZOOKEEPER_SESSION_TIMEOUT, 900 * 1000);
        return config;
    }

    public void verifyResults(String topic, int expectedCount) {
        if (topic == null) {
            topic = this.getTopicName();
        }
        if (expectedCount == -1) {
            expectedCount = sentences.length;
        }

        int foundCount = 0;
        KafkaMessageConsumer msgConsumer = null;
        try {
            msgConsumer = new KafkaMessageConsumer(getZkConnect(), topic);
            msgConsumer.consumeMessages();

            foundCount = 0;
            for (String msg : msgConsumer.getMessagesReceived()) {
                System.out.println("message: " + msg);
                if (msg.contains("cat") ||
                        msg.contains("dog") ||
                        msg.contains("bear") ||
                        msg.contains("goat") ||
                        msg.contains("SHUTDOWN")) {
                    foundCount++;
                }
            }
        } catch (Exception e) {
            e.printStackTrace();
        }

        if (foundCount != expectedCount) {
            if (msgConsumer != null) {
                System.out.println("Did not receive expected messages. Got: " +
                        msgConsumer.getMessagesReceived());
            }

            throw new RuntimeException(">>>>>>>>>>>>>>>>>>>>  Did not receive expected messages");
        }
    }

    protected void submitTopology() {

        final Config conf = getDebugConfigForStormTopology();

        cluster.submitTopology(topologyName, conf, createTopology());
    }

    public String getTopicName() {
        return topicName;
    }
}
TOP

Related Classes of AbstractStormWithKafkaTest

TOP
Copyright © 2018 www.massapi.com. All rights reserved.
All source code are property of their respective owners. Java is a trademark of Sun Microsystems, Inc and owned by ORACLE Inc. Contact coftware#gmail.com.