content: /*
   * Licensed to the Apache Software Foundation (ASF) under one or more
   * contributor license agreements.  See the NOTICE file distributed with
   * this work for additional information regarding copyright ownership.
   * The ASF 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 org.apache.activemq.broker;
  
  import javax.jms.DeliveryMode;
  import javax.jms.MessageNotWriteableException;
  import java.io.ByteArrayOutputStream;
  import java.io.DataOutputStream;
  import java.io.File;
  import java.io.IOException;
  import java.net.URI;
  import java.util.ArrayList;
  import java.util.Iterator;
  import java.util.List;
  import java.util.concurrent.TimeUnit;
  
  import org.apache.activemq.CombinationTestSupport;
  import org.apache.activemq.broker.region.RegionBroker;
  import org.apache.activemq.broker.region.policy.FixedCountSubscriptionRecoveryPolicy;
  import org.apache.activemq.broker.region.policy.PolicyEntry;
  import org.apache.activemq.broker.region.policy.PolicyMap;
  import org.apache.activemq.broker.region.policy.RoundRobinDispatchPolicy;
  import org.apache.activemq.command.ActiveMQDestination;
  import org.apache.activemq.command.ActiveMQTextMessage;
  import org.apache.activemq.command.ConnectionId;
  import org.apache.activemq.command.ConnectionInfo;
  import org.apache.activemq.command.ConsumerInfo;
  import org.apache.activemq.command.DestinationInfo;
  import org.apache.activemq.command.LocalTransactionId;
  import org.apache.activemq.command.Message;
  import org.apache.activemq.command.MessageAck;
  import org.apache.activemq.command.MessageDispatch;
  import org.apache.activemq.command.MessageId;
  import org.apache.activemq.command.ProducerInfo;
  import org.apache.activemq.command.RemoveInfo;
  import org.apache.activemq.command.SessionInfo;
  import org.apache.activemq.command.TransactionId;
  import org.apache.activemq.command.TransactionInfo;
  import org.apache.activemq.command.XATransactionId;
  import org.apache.activemq.store.PersistenceAdapter;
  import org.apache.activemq.usage.SystemUsage;
  
  public class BrokerTestSupport extends CombinationTestSupport {
  
     /**
      * Setting this to false makes the test run faster but they may be less accurate.
      */
     public static final boolean FAST_NO_MESSAGE_LEFT_ASSERT = System.getProperty("FAST_NO_MESSAGE_LEFT_ASSERT", "true").equals("true");
  
     protected RegionBroker regionBroker;
     public BrokerService broker;
     protected long idGenerator;
     protected int msgIdGenerator;
     protected int txGenerator;
     protected int tempDestGenerator;
     public PersistenceAdapter persistenceAdapter;
  
     protected String queueName = "TEST";
  
     protected int maxWait = 10000;
  
     protected SystemUsage memoryManager;
     protected PolicyMap policyMap = new PolicyMap();
  
     @Override
     protected void setUp() throws Exception {
        super.setUp();
        broker = createBroker();
        policyMap.setDefaultEntry(getDefaultPolicy());
        broker.setDestinationPolicy(policyMap);
        broker.start();
     }
  
     protected PolicyEntry getDefaultPolicy() {
        PolicyEntry policy = new PolicyEntry();
        policy.setDispatchPolicy(new RoundRobinDispatchPolicy());
        policy.setSubscriptionRecoveryPolicy(new FixedCountSubscriptionRecoveryPolicy());
        return policy;
     }
  
     protected BrokerService createBroker() throws Exception {
        BrokerService broker = BrokerFactory.createBroker(new URI("broker:()/localhost?persistent=false"));
        return broker;
     }
  
     @Override
     protected void tearDown() throws Exception {
        broker.stop();
        broker.waitUntilStopped();
        broker = null;
        regionBroker = null;
        persistenceAdapter = null;
        memoryManager = null;
        super.tearDown();
     }
  
     protected ConsumerInfo createConsumerInfo(SessionInfo sessionInfo,
                                               ActiveMQDestination destination) throws Exception {
        ConsumerInfo info = new ConsumerInfo(sessionInfo, ++idGenerator);
        info.setBrowser(false);
        info.setDestination(destination);
        info.setPrefetchSize(1000);
        info.setDispatchAsync(false);
        return info;
     }
  
     protected RemoveInfo closeConsumerInfo(ConsumerInfo consumerInfo) {
        return consumerInfo.createRemoveCommand();
     }
  
     protected ProducerInfo createProducerInfo(SessionInfo sessionInfo) throws Exception {
        ProducerInfo info = new ProducerInfo(sessionInfo, ++idGenerator);
        return info;
     }
  
     protected SessionInfo createSessionInfo(ConnectionInfo connectionInfo) throws Exception {
        SessionInfo info = new SessionInfo(connectionInfo, ++idGenerator);
        return info;
     }
  
     protected ConnectionInfo createConnectionInfo() throws Exception {
        ConnectionInfo info = new ConnectionInfo();
        info.setConnectionId(new ConnectionId("connection:" + (++idGenerator)));
        info.setClientId(info.getConnectionId().getValue());
        return info;
     }
  
     protected Message createMessage(ProducerInfo producerInfo, ActiveMQDestination destination) {
        ActiveMQTextMessage message = new ActiveMQTextMessage();
        message.setMessageId(new MessageId(producerInfo, ++msgIdGenerator));
        message.setDestination(destination);
        message.setPersistent(false);
        try {
           message.setText("Test Message Payload.");
        } catch (MessageNotWriteableException e) {
        }
        return message;
     }
  
     protected MessageAck createAck(ConsumerInfo consumerInfo, Message msg, int count, byte ackType) {
        MessageAck ack = new MessageAck();
        ack.setAckType(ackType);
        ack.setConsumerId(consumerInfo.getConsumerId());
        ack.setDestination(msg.getDestination());
        ack.setLastMessageId(msg.getMessageId());
        ack.setMessageCount(count);
        return ack;
     }
  
     protected void gc() {
        regionBroker.gc();
     }
  
     protected void profilerPause(String prompt) throws IOException {
        if (System.getProperty("profiler") != null) {
           System.out.println();
           System.out.println(prompt + "> Press enter to continue: ");
           while (System.in.read() != '\n') {
           }
           System.out.println(prompt + "> Done.");
        }
     }
  
     protected RemoveInfo closeConnectionInfo(ConnectionInfo info) {
        return info.createRemoveCommand();
     }
  
     protected RemoveInfo closeSessionInfo(SessionInfo info) {
        return info.createRemoveCommand();
     }
  
     protected RemoveInfo closeProducerInfo(ProducerInfo info) {
        return info.createRemoveCommand();
     }
  
     protected Message createMessage(ProducerInfo producerInfo, ActiveMQDestination destination, int deliveryMode) {
        Message message = createMessage(producerInfo, destination);
        message.setPersistent(deliveryMode == DeliveryMode.PERSISTENT);
        return message;
     }
  
     protected LocalTransactionId createLocalTransaction(SessionInfo info) {
        LocalTransactionId id = new LocalTransactionId(info.getSessionId().getParentId(), ++txGenerator);
        return id;
     }
  
     protected XATransactionId createXATransaction(SessionInfo info) throws IOException {
        ByteArrayOutputStream baos = new ByteArrayOutputStream();
        DataOutputStream os = new DataOutputStream(baos);
        os.writeLong(++txGenerator);
        os.close();
        byte[] bs = baos.toByteArray();
  
        XATransactionId xid = new XATransactionId();
        xid.setBranchQualifier(bs);
        xid.setGlobalTransactionId(bs);
        xid.setFormatId(55);
        return xid;
     }
  
     protected TransactionInfo createBeginTransaction(ConnectionInfo connectionInfo, TransactionId txid) {
        TransactionInfo info = new TransactionInfo(connectionInfo.getConnectionId(), txid, TransactionInfo.BEGIN);
        return info;
     }
  
     protected TransactionInfo createPrepareTransaction(ConnectionInfo connectionInfo, TransactionId txid) {
        TransactionInfo info = new TransactionInfo(connectionInfo.getConnectionId(), txid, TransactionInfo.PREPARE);
        return info;
     }
  
     protected TransactionInfo createCommitTransaction1Phase(ConnectionInfo connectionInfo, TransactionId txid) {
        TransactionInfo info = new TransactionInfo(connectionInfo.getConnectionId(), txid, TransactionInfo.COMMIT_ONE_PHASE);
        return info;
     }
  
     protected TransactionInfo createCommitTransaction2Phase(ConnectionInfo connectionInfo, TransactionId txid) {
        TransactionInfo info = new TransactionInfo(connectionInfo.getConnectionId(), txid, TransactionInfo.COMMIT_TWO_PHASE);
        return info;
     }
  
     protected TransactionInfo createRollbackTransaction(ConnectionInfo connectionInfo, TransactionId txid) {
        TransactionInfo info = new TransactionInfo(connectionInfo.getConnectionId(), txid, TransactionInfo.ROLLBACK);
        return info;
     }
  
     protected int countMessagesInQueue(StubConnection connection,
                                        ConnectionInfo connectionInfo,
                                        ActiveMQDestination destination) throws Exception {
  
        SessionInfo sessionInfo = createSessionInfo(connectionInfo);
        connection.send(sessionInfo);
        ConsumerInfo consumerInfo = createConsumerInfo(sessionInfo, destination);
        consumerInfo.setPrefetchSize(1);
        consumerInfo.setBrowser(true);
        connection.send(consumerInfo);
  
        List<Object> skipped = new ArrayList<>();
  
        // Now get the messages.
        Object m = connection.getDispatchQueue().poll(maxWait, TimeUnit.MILLISECONDS);
        int i = 0;
        while (m != null) {
           if (m instanceof MessageDispatch && ((MessageDispatch) m).getConsumerId().equals(consumerInfo.getConsumerId())) {
              MessageDispatch md = (MessageDispatch) m;
              if (md.getMessage() != null) {
                 i++;
                 connection.send(createAck(consumerInfo, md.getMessage(), 1, MessageAck.STANDARD_ACK_TYPE));
              } else {
                 break;
              }
           } else {
              skipped.add(m);
           }
           m = connection.getDispatchQueue().poll(maxWait, TimeUnit.MILLISECONDS);
        }
  
        for (Iterator<Object> iter = skipped.iterator(); iter.hasNext(); ) {
           connection.getDispatchQueue().put(iter.next());
        }
  
        connection.send(closeSessionInfo(sessionInfo));
        return i;
  
     }
  
     protected DestinationInfo createTempDestinationInfo(ConnectionInfo connectionInfo, byte destinationType) {
        DestinationInfo info = new DestinationInfo();
        info.setConnectionId(connectionInfo.getConnectionId());
        info.setOperationType(DestinationInfo.ADD_OPERATION_TYPE);
        info.setDestination(ActiveMQDestination.createDestination(info.getConnectionId() + ":" + (++tempDestGenerator), destinationType));
        return info;
     }
  
     protected ActiveMQDestination createDestinationInfo(StubConnection connection,
                                                         ConnectionInfo connectionInfo1,
                                                         byte destinationType) throws Exception {
        if ((destinationType & ActiveMQDestination.TEMP_MASK) != 0) {
           DestinationInfo info = createTempDestinationInfo(connectionInfo1, destinationType);
           connection.send(info);
           return info.getDestination();
        } else {
           return ActiveMQDestination.createDestination(queueName, destinationType);
        }
     }
  
     protected DestinationInfo closeDestinationInfo(DestinationInfo info) {
        info.setOperationType(DestinationInfo.REMOVE_OPERATION_TYPE);
        info.setTimeout(0);
        return info;
     }
  
     public static void recursiveDelete(File f) {
        if (f.isDirectory()) {
           File[] files = f.listFiles();
           for (int i = 0; i < files.length; i++) {
              recursiveDelete(files[i]);
           }
        }
        f.delete();
     }
  
     protected StubConnection createConnection() throws Exception {
        return new StubConnection(broker);
     }
  
     public Message receiveMessage(StubConnection connection) throws InterruptedException {
        return receiveMessage(connection, maxWait);
     }
  
     public Message receiveMessage(StubConnection connection, long timeout) throws InterruptedException {
        while (true) {
           Object o = connection.getDispatchQueue().poll(timeout, TimeUnit.MILLISECONDS);
  
           if (o == null) {
              return null;
           }
           if (o instanceof MessageDispatch) {
  
              MessageDispatch dispatch = (MessageDispatch) o;
              if (dispatch.getMessage() == null) {
                 return null;
              }
              dispatch.setMessage(dispatch.getMessage().copy());
              dispatch.getMessage().setRedeliveryCounter(dispatch.getRedeliveryCounter());
              return dispatch.getMessage();
           }
        }
     }
  
     protected void assertNoMessagesLeft(StubConnection connection) throws InterruptedException {
        long wait = FAST_NO_MESSAGE_LEFT_ASSERT ? 0 : maxWait;
        while (true) {
           Object o = connection.getDispatchQueue().poll(wait, TimeUnit.MILLISECONDS);
           if (o == null) {
              return;
           }
           if (o instanceof MessageDispatch && ((MessageDispatch) o).getMessage() != null) {
              fail("Received a message: " + ((MessageDispatch) o).getMessage().getMessageId());
           }
        }
     }
  }
  