001/*
002 * Licensed to the Apache Software Foundation (ASF) under one or more
003 * contributor license agreements.  See the NOTICE file distributed with
004 * this work for additional information regarding copyright ownership.
005 * The ASF licenses this file to You under the Apache License, Version 2.0
006 * (the "License"); you may not use this file except in compliance with
007 * the License.  You may obtain a copy of the License at
008 *
009 *      http://www.apache.org/licenses/LICENSE-2.0
010 *
011 * Unless required by applicable law or agreed to in writing, software
012 * distributed under the License is distributed on an "AS IS" BASIS,
013 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
014 * See the License for the specific language governing permissions and
015 * limitations under the License.
016 */
017package org.apache.activemq.junit;
018
019import java.net.URI;
020
021import org.apache.activemq.ActiveMQConnectionFactory;
022import org.apache.activemq.broker.BrokerFactory;
023import org.apache.activemq.broker.BrokerPlugin;
024import org.apache.activemq.broker.BrokerService;
025import org.apache.activemq.broker.region.Destination;
026import org.apache.activemq.broker.region.Queue;
027import org.apache.activemq.broker.region.Topic;
028import org.apache.activemq.broker.region.policy.PolicyEntry;
029import org.apache.activemq.broker.region.policy.PolicyMap;
030import org.apache.activemq.plugin.StatisticsBrokerPlugin;
031import org.apache.activemq.pool.PooledConnectionFactory;
032import org.junit.rules.ExternalResource;
033import org.slf4j.Logger;
034import org.slf4j.LoggerFactory;
035
036/**
037 * A JUnit Rule that embeds an ActiveMQ broker into a test.
038 */
039public class EmbeddedActiveMQBroker extends ExternalResource {
040    Logger log = LoggerFactory.getLogger(this.getClass());
041
042    BrokerService brokerService;
043
044    /**
045     * Create an embedded ActiveMQ broker using defaults
046     *
047     * The defaults are:
048     *  - the broker name is 'embedded-broker'
049     *  - JMX is enabled but no management connector will be created.
050     *  - Persistence is disabled
051     *
052     */
053    public EmbeddedActiveMQBroker() {
054        brokerService = new BrokerService();
055        brokerService.setUseJmx(true);
056        brokerService.getManagementContext().setCreateConnector(false);
057        brokerService.setUseShutdownHook(false);
058        brokerService.setPersistent(false);
059        brokerService.setBrokerName("embedded-broker");
060    }
061
062    /**
063     * Create an embedded ActiveMQ broker using a configuration URI
064     */
065    public EmbeddedActiveMQBroker(String configurationURI ) {
066        try {
067            brokerService = BrokerFactory.createBroker(configurationURI);
068        } catch (Exception ex) {
069            throw new RuntimeException("Exception encountered creating embedded ActiveMQ broker from configuration URI: " + configurationURI, ex);
070        }
071    }
072
073    /**
074     * Create an embedded ActiveMQ broker using a configuration URI
075     */
076    public EmbeddedActiveMQBroker(URI configurationURI ) {
077        try {
078            brokerService = BrokerFactory.createBroker(configurationURI);
079        } catch (Exception ex) {
080            throw new RuntimeException("Exception encountered creating embedded ActiveMQ broker from configuration URI: " + configurationURI, ex);
081        }
082    }
083
084    /**
085     * Customize the configuration of the embedded ActiveMQ broker
086     *
087     * This method is called before the embedded ActiveMQ broker is started, and can
088     * be overridden to this method to customize the broker configuration.
089     */
090    protected void configure() {}
091
092    /**
093     * Start the embedded ActiveMQ broker, blocking until the broker has successfully started.
094     * <p/>
095     * The broker will normally be started by JUnit using the before() method.  This method allows the broker to
096     * be started manually to support advanced testing scenarios.
097     */
098    public void start() {
099        try {
100            this.configure();
101            brokerService.start();
102        } catch (Exception ex) {
103            throw new RuntimeException("Exception encountered starting embedded ActiveMQ broker: {}" + this.getBrokerName(), ex);
104        }
105
106        brokerService.waitUntilStarted();
107    }
108
109    /**
110     * Stop the embedded ActiveMQ broker, blocking until the broker has stopped.
111     * <p/>
112     * The broker will normally be stopped by JUnit using the after() method.  This method allows the broker to
113     * be stopped manually to support advanced testing scenarios.
114     */
115    public void stop() {
116        if (!brokerService.isStopped()) {
117            try {
118                brokerService.stop();
119            } catch (Exception ex) {
120                log.warn("Exception encountered stopping embedded ActiveMQ broker: {}" + this.getBrokerName(), ex);
121            }
122        }
123
124        brokerService.waitUntilStopped();
125    }
126
127    /**
128     * Start the embedded ActiveMQ Broker
129     * <p/>
130     * Invoked by JUnit to setup the resource
131     */
132    @Override
133    protected void before() throws Throwable {
134        log.info("Starting embedded ActiveMQ broker: {}", this.getBrokerName());
135
136        this.start();
137
138        super.before();
139    }
140
141    /**
142     * Stop the embedded ActiveMQ Broker
143     * <p/>
144     * Invoked by JUnit to tear down the resource
145     */
146    @Override
147    protected void after() {
148        log.info("Stopping Embedded ActiveMQ Broker: {}", this.getBrokerName());
149
150        super.after();
151
152        this.stop();
153    }
154
155    /**
156     * Create an ActiveMQConnectionFactory for the embedded ActiveMQ Broker
157     *
158     * @return a new ActiveMQConnectionFactory
159     */
160    public ActiveMQConnectionFactory createConnectionFactory() {
161        ActiveMQConnectionFactory connectionFactory = new ActiveMQConnectionFactory();
162        connectionFactory.setBrokerURL(brokerService.getVmConnectorURI().toString());
163        return connectionFactory;
164    }
165
166    /**
167     * Create an PooledConnectionFactory for the embedded ActiveMQ Broker
168     *
169     * @return a new PooledConnectionFactory
170     */
171    public PooledConnectionFactory createPooledConnectionFactory() {
172        ActiveMQConnectionFactory connectionFactory = createConnectionFactory();
173
174        PooledConnectionFactory pooledConnectionFactory = new PooledConnectionFactory(connectionFactory);
175
176        return pooledConnectionFactory;
177    }
178
179    /**
180     * Get the BrokerService for the embedded ActiveMQ broker.
181     * <p/>
182     * This may be required for advanced configuration of the BrokerService.
183     *
184     * @return the embedded ActiveMQ broker
185     */
186    public BrokerService getBrokerService() {
187        return brokerService;
188    }
189
190    /**
191     * Get the VM URL for the embedded ActiveMQ Broker
192     * <p/>
193     * NOTE:  The option is precreate=false option is appended to the URL to avoid the automatic creation of brokers
194     * and the resulting duplicate broker errors
195     *
196     * @return the VM URL for the embedded broker
197     */
198    public String getVmURL() {
199        return String.format("failover:(%s?create=false)", brokerService.getVmConnectorURI().toString());
200    }
201
202    /**
203     * Get the name of the embedded ActiveMQ Broker
204     *
205     * @return name of the embedded broker
206     */
207    public String getBrokerName() {
208        return brokerService.getBrokerName();
209    }
210
211    public void setBrokerName(String brokerName) {
212        brokerService.setBrokerName(brokerName);
213    }
214
215    public boolean isStatisticsPluginEnabled() {
216        BrokerPlugin[] plugins = brokerService.getPlugins();
217
218        if (null != plugins) {
219            for (BrokerPlugin plugin : plugins) {
220                if (plugin instanceof StatisticsBrokerPlugin) {
221                    return true;
222                }
223            }
224        }
225
226        return false;
227    }
228
229    public void enableStatisticsPlugin() {
230        if (!isStatisticsPluginEnabled()) {
231            BrokerPlugin[] newPlugins;
232            BrokerPlugin[] currentPlugins = brokerService.getPlugins();
233            if (null != currentPlugins && 0 < currentPlugins.length) {
234                newPlugins = new BrokerPlugin[currentPlugins.length + 1];
235
236                System.arraycopy(currentPlugins, 0, newPlugins, 0, currentPlugins.length);
237            } else {
238                newPlugins = new BrokerPlugin[1];
239            }
240
241            newPlugins[newPlugins.length - 1] = new StatisticsBrokerPlugin();
242
243            brokerService.setPlugins(newPlugins);
244        }
245    }
246
247    public void disableStatisticsPlugin() {
248        if (isStatisticsPluginEnabled()) {
249            BrokerPlugin[] currentPlugins = brokerService.getPlugins();
250            if (1 < currentPlugins.length) {
251                BrokerPlugin[] newPlugins = new BrokerPlugin[currentPlugins.length - 1];
252
253                int i = 0;
254                for (BrokerPlugin plugin : currentPlugins) {
255                    if (!(plugin instanceof StatisticsBrokerPlugin)) {
256                        newPlugins[i++] = plugin;
257                    }
258                }
259                brokerService.setPlugins(newPlugins);
260            } else {
261                brokerService.setPlugins(null);
262            }
263
264        }
265    }
266
267    public boolean isAdvisoryForDeliveryEnabled() {
268        return getDefaultPolicyEntry().isAdvisoryForDelivery();
269    }
270
271    public void enableAdvisoryForDelivery() {
272        getDefaultPolicyEntry().setAdvisoryForDelivery(true);
273    }
274
275    public void disableAdvisoryForDelivery() {
276        getDefaultPolicyEntry().setAdvisoryForDelivery(false);
277    }
278
279    public boolean isAdvisoryForConsumedEnabled() {
280        return getDefaultPolicyEntry().isAdvisoryForConsumed();
281    }
282
283    public void enableAdvisoryForConsumed() {
284        getDefaultPolicyEntry().setAdvisoryForConsumed(true);
285    }
286
287    public void disableAdvisoryForConsumed() {
288        getDefaultPolicyEntry().setAdvisoryForConsumed(false);
289    }
290
291    public boolean isAdvisoryForDiscardingMessagesEnabled() {
292        return getDefaultPolicyEntry().isAdvisoryForDiscardingMessages();
293    }
294
295    public void enableAdvisoryForDiscardingMessages() {
296        getDefaultPolicyEntry().setAdvisoryForDiscardingMessages(true);
297    }
298
299    public void disableAdvisoryForDiscardingMessages() {
300        getDefaultPolicyEntry().setAdvisoryForDiscardingMessages(false);
301    }
302
303    public boolean isAdvisoryForFastProducersEnabled() {
304        return getDefaultPolicyEntry().isAdvisoryForFastProducers();
305    }
306
307    public void enableAdvisoryForFastProducers() {
308        getDefaultPolicyEntry().setAdvisoryForFastProducers(true);
309    }
310
311    public void disableAdvisoryForFastProducers() {
312        getDefaultPolicyEntry().setAdvisoryForFastProducers(false);
313    }
314
315    public boolean isAdvisoryForSlowConsumersEnabled() {
316        return getDefaultPolicyEntry().isAdvisoryForSlowConsumers();
317    }
318
319    public void enableAdvisoryForSlowConsumers() {
320        getDefaultPolicyEntry().setAdvisoryForSlowConsumers(true);
321    }
322
323    public void disableAdvisoryForSlowConsumers() {
324        getDefaultPolicyEntry().setAdvisoryForSlowConsumers(false);
325    }
326
327    /**
328     * Get the number of messages in a specific JMS Destination.
329     * <p/>
330     * The full name of the JMS destination including the prefix should be provided - i.e. queue:myQueue
331     * or topic:myTopic.  If the destination type prefix is not included in the destination name, a prefix
332     * of "queue:" is assumed.
333     *
334     * @param fullDestinationName the full name of the JMS Destination
335     * @return the number of messages in the JMS Destination
336     */
337    public int getMessageCount(String fullDestinationName) throws Exception {
338        final int QUEUE_TYPE = 1;
339        final int TOPIC_TYPE = 2;
340
341        if (null == brokerService) {
342            throw new IllegalStateException("BrokerService has not yet been created - was before() called?");
343        }
344
345        int destinationType = QUEUE_TYPE;
346        String destinationName = fullDestinationName;
347
348        if (fullDestinationName.startsWith("queue:")) {
349            destinationName = fullDestinationName.substring(fullDestinationName.indexOf(':') + 1);
350        } else if (fullDestinationName.startsWith("topic:")) {
351            destinationType = TOPIC_TYPE;
352            destinationName = fullDestinationName.substring(fullDestinationName.indexOf(':') + 1);
353        }
354
355        int messageCount = -1;
356        boolean foundDestination = false;
357        for (Destination destination : brokerService.getBroker().getDestinationMap().values()) {
358            String tmpName = destination.getName();
359            if (tmpName.equalsIgnoreCase(destinationName)) {
360                switch (destinationType) {
361                    case QUEUE_TYPE:
362                        if (destination instanceof Queue) {
363                            messageCount = destination.getMessageStore().getMessageCount();
364                            foundDestination = true;
365                        }
366                        break;
367                    case TOPIC_TYPE:
368                        if (destination instanceof Topic) {
369                            messageCount = destination.getMessageStore().getMessageCount();
370                            foundDestination = true;
371                        }
372                        break;
373                    default:
374                        // Should never see this
375                        log.error("Type didn't match: {}", destination.getClass().getName());
376                }
377            }
378            if (foundDestination) {
379                break;
380            }
381        }
382
383        if (!foundDestination) {
384            log.warn("Didn't find destination {} in broker {}", fullDestinationName, getBrokerName());
385        }
386
387        return messageCount;
388    }
389
390    private PolicyEntry getDefaultPolicyEntry() {
391        PolicyMap destinationPolicy = brokerService.getDestinationPolicy();
392        if (null == destinationPolicy) {
393            destinationPolicy = new PolicyMap();
394            brokerService.setDestinationPolicy(destinationPolicy);
395        }
396
397        PolicyEntry defaultEntry = destinationPolicy.getDefaultEntry();
398        if (null == defaultEntry) {
399            defaultEntry = new PolicyEntry();
400            destinationPolicy.setDefaultEntry(defaultEntry);
401        }
402
403        return defaultEntry;
404    }
405}