xyz.cloudbans.barker.amqp.AmqpBarkerTest.java Source code

Introduction

Here is the source code for xyz.cloudbans.barker.amqp.AmqpBarkerTest.java

Source

/*
 * Copyright 2016 CloudBans (https://cloudbans.xyz)
 *
 *    Licensed 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 xyz.cloudbans.barker.amqp;

import com.rabbitmq.client.Channel;
import com.rabbitmq.client.Connection;
import com.rabbitmq.client.ConnectionFactory;
import org.hamcrest.CoreMatchers;
import org.junit.AfterClass;
import org.junit.Assert;
import org.junit.Assume;
import org.junit.BeforeClass;
import org.junit.Test;
import xyz.cloudbans.barker.Barker;
import xyz.cloudbans.barker.Message;
import xyz.cloudbans.barker.Subscription;
import xyz.cloudbans.barker.listener.BarkerListener;
import xyz.cloudbans.barker.listener.OneTimeListener;

import java.io.IOException;
import java.util.UUID;
import java.util.concurrent.CountDownLatch;
import java.util.concurrent.atomic.AtomicBoolean;
import java.util.concurrent.atomic.AtomicInteger;
import java.util.logging.Logger;

public class AmqpBarkerTest {

    private static final Logger LOGGER = Logger.getLogger(AmqpBarkerTest.class.getName());
    private static final String TOPIC_NAME = "test";

    private static Connection connection;

    @BeforeClass
    public static void prepare() throws Throwable {
        String amqpUri = System.getenv("AMQP_URI");
        Assume.assumeNotNull("AMQP_URI is not set.", amqpUri, CoreMatchers.notNullValue());

        ConnectionFactory connectionFactory = new ConnectionFactory();
        connectionFactory.setUri(amqpUri);
        try {
            connection = connectionFactory.newConnection();
        } catch (IOException e) {
            Assume.assumeNoException("Connection to AMQP broker cannot be established.", e);
        }
    }

    @Test
    public void testConsume() throws Throwable {
        Channel channel = openChannel();
        AmqpBarkerOptions options = createOptions();
        Barker barker = new AmqpBarker(channel, options);
        final CountDownLatch countDownLatch = new CountDownLatch(4);
        final AtomicInteger counter = new AtomicInteger();
        final AtomicBoolean failed = new AtomicBoolean();

        barker.subscribe(TOPIC_NAME, new BarkerListener() {
            @Override
            public void onMessage(Subscription subscription, Message message) {
                if (!"test".equals(subscription.getTopic()) || !new String(message.getBody()).startsWith("Test"))
                    failed.set(true);
                countDownLatch.countDown();
                counter.incrementAndGet();
                message.acknowledge();
            }
        });

        for (int i = 0; i < 4; i++)
            barker.submit(TOPIC_NAME, ("Test #" + i).getBytes());

        countDownLatch.await();
        // Wait  a short moment for possible more messages
        Thread.sleep(100);

        Assert.assertFalse(failed.get());
        Assert.assertEquals(4, counter.get());

        cleanup(options, channel);
        channel.abort();
    }

    @Test
    public void testOneTimeConsume() throws Throwable {
        Channel channel = openChannel();
        AmqpBarkerOptions options = createOptions();
        Barker barker = new AmqpBarker(channel, options);
        final CountDownLatch countDownLatch = new CountDownLatch(1);
        final AtomicInteger counter = new AtomicInteger();
        final AtomicBoolean failed = new AtomicBoolean();

        barker.subscribe(TOPIC_NAME, new OneTimeListener() {
            @Override
            protected boolean handleMessage(Subscription subscription, Message message) {
                if (!"test".equals(subscription.getTopic()) || !new String(message.getBody()).equals("Test"))
                    failed.set(true);
                countDownLatch.countDown();
                counter.incrementAndGet();
                return true;
            }
        });

        for (int i = 0; i < 2; i++)
            barker.submit(TOPIC_NAME, "Test".getBytes());

        countDownLatch.await();
        // Wait  a short moment for possible more messages
        Thread.sleep(100);

        Assert.assertFalse(failed.get());
        Assert.assertEquals(1, counter.get());

        cleanup(options, channel);
        channel.abort();
    }

    @AfterClass
    public static void shutdown() throws Throwable {
        if (connection != null)
            connection.abort();
    }

    private Channel openChannel() throws IOException {
        Channel channel = connection.createChannel();
        channel.basicQos(1);
        return channel;
    }

    private AmqpBarkerOptions createOptions() {
        UUID uuid = UUID.randomUUID();
        return new AmqpBarkerOptions().setInputExchange("barker-test#" + uuid).setQueuePrefix("test#" + uuid + ".");
    }

    private void cleanup(AmqpBarkerOptions options, Channel channel) {
        try {
            channel.exchangeDelete(options.getInputExchange());
        } catch (IOException ignore) {
            LOGGER.info("An non-critical exception occurred while deleting exchange.");
        }
        try {
            channel.queueDelete(options.getQueuePrefix() + TOPIC_NAME);
        } catch (IOException ignore) {
            LOGGER.info("An non-critical exception occurred while deleting queue.");
        }
    }
}