package com.walmartlabs.mupd8;

import java.io.File;
import com.google.common.io.Files;
import com.walmartlabs.mupd8.application.Mupd8DataPair;
import junit.framework.TestCase;
import kafka.producer.KeyedMessage;
import kafka.producer.ProducerConfig;
import kafka.server.KafkaConfig;
import kafka.javaapi.producer.Producer;

import java.util.List;
import java.util.ArrayList;
import java.util.Properties;

import kafka.server.KafkaServerStartable;
import org.apache.curator.test.TestingServer;

public class KafkaSourceTest extends TestCase {
    private int port = 9092;
    private int brokerId = 1;
    private TestingServer zkServer;
    private KafkaServerStartable kafkaServer;
    private KafkaSource kafkaSource;
    private File logDir;
    private String topic = "mupd8-test-topic";
    private String key = "Kf1";
    private String val = "data1";
    private String mesg = "{ \"" + key + "\" : \"" + val + "\" }";

    public void setUp() throws Exception {
        zkServer = new TestingServer();
        logDir = Files.createTempDir();
        kafkaServer = startKafkaServer(zkServer.getConnectString(), logDir.getAbsolutePath());
        produceJsonMessage();
        kafkaSource = getKafkaSource();
    }

    public void tearDown() throws Exception {
        kafkaSource.closeSource();
        kafkaServer.shutdown();
        zkServer.stop();
        logDir.delete();
    }

    public void testHasNext()  {
        assertTrue("HasNext should return true when data", kafkaSource.hasNext());
        assertTrue("HasNext should not consume data",kafkaSource.hasNext());
    }

    public void testNext() {
        assertTrue("HasNext should return true when data", kafkaSource.hasNext());
        Mupd8DataPair mupd8DataPair = kafkaSource.getNextDataPair();
        assertEquals("Next should return the correct message",val,mupd8DataPair._key.toString());
        produceNonJsonMessage();
        mupd8DataPair = kafkaSource.getNextDataPair();
        assertEquals("Next should return the next correct message",val, new String(mupd8DataPair._value));
        assertNull("Next should return the null val for missing key", mupd8DataPair._key);
    }

    private KafkaSource getKafkaSource()  {
        List<String> argsList = new ArrayList<String>();
        argsList.add(zkServer.getConnectString());
        argsList.add("test-mupd8-consumer");
        argsList.add(topic);
        argsList.add(key);
        return new KafkaSource(argsList);
    }

    private KafkaServerStartable startKafkaServer(String zkConnect, String logDir){
        Properties props = new Properties();
        props.put("port", "9092");
        props.put("broker.id", "1");
        props.put("log.dir", logDir);
        props.put("zookeeper.connect",zkConnect);
        KafkaConfig kafkaConfig = new KafkaConfig(props);
        KafkaServerStartable server = new KafkaServerStartable(kafkaConfig);
        server.startup();
        return server;
    }

    private void produceJsonMessage(){
        produceMessage(topic,mesg);
    }

    private void produceNonJsonMessage(){
        produceMessage(topic,val);
    }

    //Produces messages and also creates topic if not already created
    private void produceMessage(String topic, String mesg){
        Properties props = new Properties();
        props.put("metadata.broker.list", "localhost:9092");
        props.put("serializer.class", "kafka.serializer.StringEncoder");
        props.put("producer.type", "async");
        props.put("batch.size", "1");
        ProducerConfig config = new ProducerConfig(props);
        Producer producer = new Producer<String, String>(config);
        producer.send(new KeyedMessage(topic,mesg));
    }

}

@@ -57,91 +57,81 @@
 ort 
-com.google.common.io.Files;%0Aimport com.walmartlabs.mupd8.application.Mupd8DataPair;
+java.util.ArrayList;%0Aimport java.util.List;%0Aimport java.util.Properties;%0A
 %0Aimp
@@ -177,43 +177,15 @@
 fka.
-producer.KeyedMessage;%0Aimport kafka
+javaapi
 .pro
@@ -198,22 +198,16 @@
 Producer
-Config
 ;%0Aimport
@@ -217,26 +217,29 @@
 fka.
-server.KafkaConfig
+producer.KeyedMessage
 ;%0Aim
@@ -245,32 +245,24 @@
 mport kafka.
-javaapi.
 producer.Pro
@@ -270,91 +270,48 @@
 ucer
-;%0A%0Aimport java.util.List;%0Aimport java.util.ArrayList;%0Aimport java.util.Properties;%0A
+Config;%0Aimport kafka.server.KafkaConfig;
 %0Aimp
@@ -345,24 +345,25 @@
 rStartable;%0A
+%0A
 import org.a
@@ -401,114 +401,147 @@
 r;%0A%0A
-public class KafkaSourceTest extends TestCase %7B%0A    private int port = 9092;%0A    private int brokerId = 1;
+import com.google.common.io.Files;%0Aimport com.walmartlabs.mupd8.application.Mupd8DataPair;%0A%0Apublic class KafkaSourceTest extends TestCase %7B
 %0A   
@@ -862,16 +862,30 @@
 %5C%22 %7D%22;%0A%0A
+    @Override%0A
     publ
@@ -1163,32 +1163,46 @@
 ource();%0A    %7D%0A%0A
+    @Override%0A
     public void 
@@ -3542,16 +3542,32 @@
 Producer
+%3CString, String%3E
  produce
@@ -3642,24 +3642,40 @@
 KeyedMessage
+%3CString, String%3E
 (topic,mesg)

Fix warning in KafkaSourceTest