com.metamx.tranquility.kafka.writer.WriterControllerTest.java Source code

Java tutorial

Introduction

Here is the source code for com.metamx.tranquility.kafka.writer.WriterControllerTest.java

Source

/*
 * Licensed to Metamarkets Group Inc. (Metamarkets) under one
 * or more contributor license agreements.  See the NOTICE file
 * distributed with this work for additional information
 * regarding copyright ownership.  Metamarkets licenses this file
 * to you under the Apache License, Version 2.0 (the
 * "License"); you may not use this file except in compliance
 * with the License.  You may obtain a copy of the License at
 *
 *   http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing,
 * software distributed under the License is distributed on an
 * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
 * KIND, either express or implied.  See the License for the
 * specific language governing permissions and limitations
 * under the License.
 */

package com.metamx.tranquility.kafka.writer;

import com.fasterxml.jackson.databind.ObjectMapper;
import com.google.common.collect.ImmutableMap;
import com.metamx.tranquility.config.DataSourceConfig;
import com.metamx.tranquility.kafka.KafkaConsumerTest;
import com.metamx.tranquility.kafka.model.MessageCounters;
import com.metamx.tranquility.kafka.model.PropertiesBasedKafkaConfig;
import io.druid.query.aggregation.AggregatorFactory;
import io.druid.segment.indexing.DataSchema;
import io.druid.segment.indexing.RealtimeIOConfig;
import io.druid.segment.indexing.RealtimeTuningConfig;
import io.druid.segment.realtime.FireDepartment;
import io.druid.segment.realtime.FireDepartmentMetrics;
import io.druid.segment.realtime.firehose.LocalFirehoseFactory;
import io.druid.segment.realtime.plumber.Plumber;
import io.druid.segment.realtime.plumber.PlumberSchool;
import junit.framework.Assert;
import org.easymock.EasyMock;
import org.junit.Before;
import org.junit.Test;
import org.skife.config.ConfigurationObjectFactory;

import java.io.IOException;
import java.util.HashMap;
import java.util.Map;
import java.util.Properties;

public class WriterControllerTest {
    class TestableWriterController extends WriterController {

        private Map<String, TranquilityEventWriter> mockWriters = new HashMap<>();

        public TestableWriterController(
                Map<String, DataSourceConfig<PropertiesBasedKafkaConfig>> dataSourceConfigs) {
            super(dataSourceConfigs);
        }

        protected TranquilityEventWriter createWriter(String topic,
                DataSourceConfig<PropertiesBasedKafkaConfig> dataSourceConfig) {
            TranquilityEventWriter writer = EasyMock.mock(TranquilityEventWriter.class);
            mockWriters.put(topic, writer);
            return writer;
        }
    }

    private TestableWriterController writerController;

    @Before
    public void setUp() throws IOException {
        Properties props = new Properties();
        props.setProperty("kafka.zookeeper.connect", "localhost:2181");
        props.setProperty("zookeeper.connect", "localhost:2181");

        Properties propsTwitter = new Properties(props);
        propsTwitter.setProperty("topicPattern", "twitter");

        Properties propsTest = new Properties(props);
        propsTest.setProperty("topicPattern", "test[0-9]");
        propsTest.setProperty("useTopicAsDataSource", "true");

        FireDepartment fdTwitter = new FireDepartment(
                new DataSchema("twitter", null, new AggregatorFactory[] {}, null, new ObjectMapper()),
                new RealtimeIOConfig(new LocalFirehoseFactory(null, null, null), new PlumberSchool() {
                    @Override
                    public Plumber findPlumber(DataSchema schema, RealtimeTuningConfig config,
                            FireDepartmentMetrics metrics) {
                        return null;
                    }
                }, null), null);

        FireDepartment fdTest = new FireDepartment(
                new DataSchema("test[0-9]", null, new AggregatorFactory[] {}, null, new ObjectMapper()),
                new RealtimeIOConfig(new LocalFirehoseFactory(null, null, null), new PlumberSchool() {
                    @Override
                    public Plumber findPlumber(DataSchema schema, RealtimeTuningConfig config,
                            FireDepartmentMetrics metrics) {
                        return null;
                    }
                }, null), null);

        Map<String, DataSourceConfig<PropertiesBasedKafkaConfig>> datasourceConfigs = ImmutableMap.of("twitter",
                new DataSourceConfig<>("twitter",
                        new ConfigurationObjectFactory(propsTwitter).build(PropertiesBasedKafkaConfig.class),
                        KafkaConsumerTest.fireDepartmentToScalaMap(fdTwitter)),
                "test[0-9]",
                new DataSourceConfig<>("test[0-9]",
                        new ConfigurationObjectFactory(propsTest).build(PropertiesBasedKafkaConfig.class),
                        KafkaConsumerTest.fireDepartmentToScalaMap(fdTest)));

        writerController = new TestableWriterController(datasourceConfigs);
    }

    @Test
    public void testGetWriter() {
        Assert.assertEquals(0, writerController.mockWriters.size());

        TranquilityEventWriter writerTwitter = writerController.getWriter("twitter");

        Assert.assertEquals(1, writerController.mockWriters.size());
        Assert.assertSame(writerTwitter, writerController.mockWriters.get("twitter"));
        Assert.assertSame(writerTwitter, writerController.getWriter("twitter"));

        TranquilityEventWriter writerTest0 = writerController.getWriter("test0");
        TranquilityEventWriter writerTest1 = writerController.getWriter("test1");

        Assert.assertEquals(3, writerController.mockWriters.size());
        Assert.assertSame(writerTest0, writerController.getWriter("test0"));
        Assert.assertSame(writerTest1, writerController.getWriter("test1"));
    }

    @Test(expected = RuntimeException.class)
    public void testGetWriterInvalid() {
        writerController.getWriter("test10");
    }

    @Test
    public void testFlushAll() throws InterruptedException {
        writerController.getWriter("test0");
        writerController.getWriter("test1");
        TranquilityEventWriter mock0 = writerController.mockWriters.get("test0");
        TranquilityEventWriter mock1 = writerController.mockWriters.get("test1");

        mock0.flush();
        mock1.flush();
        EasyMock.expect(mock0.getMessageCounters()).andReturn(new MessageCounters(1, 2, 3, 0));
        EasyMock.expect(mock1.getMessageCounters()).andReturn(new MessageCounters(4, 5, 6, 0));
        EasyMock.replay(mock0, mock1);

        Map<String, MessageCounters> results = writerController.flushAll();

        EasyMock.verify(mock0, mock1);

        Assert.assertEquals(2, results.size());

        Assert.assertEquals(1, results.get("test0").getReceivedCount());
        Assert.assertEquals(2, results.get("test0").getSentCount());
        Assert.assertEquals(3, results.get("test0").getDroppedCount());
        Assert.assertEquals(0, results.get("test0").getUnparseableCount());

        Assert.assertEquals(4, results.get("test1").getReceivedCount());
        Assert.assertEquals(5, results.get("test1").getSentCount());
        Assert.assertEquals(6, results.get("test1").getDroppedCount());
        Assert.assertEquals(0, results.get("test1").getUnparseableCount());
    }

    @Test
    public void testStop() {
        writerController.getWriter("test0");
        writerController.getWriter("test1");
        TranquilityEventWriter mock0 = writerController.mockWriters.get("test0");
        TranquilityEventWriter mock1 = writerController.mockWriters.get("test1");

        mock0.stop();
        mock1.stop();
        EasyMock.replay(mock0, mock1);

        writerController.stop();

        EasyMock.verify(mock0, mock1);
    }
}