Adding support for amqp custom headers.

jlavallee · Sep 25, 2014 · 8c45acf · 8c45acf
1 parent 073ee2b
commit 8c45acf
Show file tree

Hide file tree

Showing 2 changed files with 39 additions and 0 deletions.
diff --git a/src/main/com/zeroclue/jmeter/protocol/amqp/AMQPPublisher.java b/src/main/com/zeroclue/jmeter/protocol/amqp/AMQPPublisher.java
@@ -3,11 +3,14 @@
 import com.rabbitmq.client.AMQP;
 import java.io.IOException;
 import java.security.*;
+import java.util.*;
 
 import com.rabbitmq.client.MessageProperties;
+import org.apache.jmeter.config.Arguments;
 import org.apache.jmeter.samplers.Entry;
 import org.apache.jmeter.samplers.Interruptible;
 import org.apache.jmeter.samplers.SampleResult;
+import org.apache.jmeter.testelement.property.TestElementProperty;
 import org.apache.jorphan.logging.LoggingManager;
 import org.apache.log.Logger;
 
@@ -35,6 +38,7 @@ public class AMQPPublisher extends AMQPSampler implements Interruptible {
     private final static String MESSAGE_TYPE = "AMQPPublisher.MessageType";
     private final static String REPLY_TO_QUEUE = "AMQPPublisher.ReplyToQueue";
     private final static String CORRELATION_ID = "AMQPPublisher.CorrelationId";
+    private final static String HEADERS = "AMQPPublisher.Headers";
 
     public static boolean DEFAULT_PERSISTENT = false;
     private final static String PERSISTENT = "AMQPConsumer.Persistent";
@@ -78,6 +82,7 @@ public SampleResult sample(Entry e) {
         result.sampleStart(); // Start timing
         try {
             AMQP.BasicProperties messageProperties = getProperties();
+            messageProperties.setHeaders(prepareHeaders());
             byte[] messageBytes = getMessageBytes();
 
             for (int idx = 0; idx < loop; idx++) {
@@ -176,6 +181,13 @@ public void setCorrelationId(String content) {
         setProperty(CORRELATION_ID, content);
     }
 
+    public Arguments getHeaders() {
+        return (Arguments) getProperty(HEADERS).getObjectValue();
+    }
+
+    public void setHeaders(Arguments headers) {
+        setProperty(new TestElementProperty(HEADERS, headers));
+    }
 
     public Boolean getPersistent() {
         return getPropertyAsBoolean(PERSISTENT, DEFAULT_PERSISTENT);
@@ -232,4 +244,13 @@ protected boolean initChannel() throws IOException, NoSuchAlgorithmException, Ke
         }
         return ret;
     }
+
+    private Map<String, Object> prepareHeaders() {
+        Map<String, Object> result = new HashMap<String, Object>();
+        Map<String, String> source = getHeaders().getArgumentsAsMap();
+        for (Map.Entry<String, String> item : source.entrySet()) {
+            result.put(item.getKey(), item.getValue());
+        }
+        return result;
+    }
 }
diff --git a/src/main/com/zeroclue/jmeter/protocol/amqp/gui/AMQPPublisherGui.java b/src/main/com/zeroclue/jmeter/protocol/amqp/gui/AMQPPublisherGui.java
@@ -4,6 +4,8 @@
 
 import javax.swing.*;
 
+import org.apache.jmeter.config.Arguments;
+import org.apache.jmeter.config.gui.ArgumentsPanel;
 import org.apache.jmeter.testelement.TestElement;
 import org.apache.jorphan.gui.JLabeledTextArea;
 import org.apache.jorphan.gui.JLabeledTextField;
@@ -40,6 +42,8 @@ public class AMQPPublisherGui extends AMQPSamplerGui {
     private JCheckBox persistent = new JCheckBox("Persistent?", AMQPPublisher.DEFAULT_PERSISTENT);
     private JCheckBox useTx = new JCheckBox("Use Transactions?", AMQPPublisher.DEFAULT_USE_TX);
 
+    private ArgumentsPanel headers = new ArgumentsPanel("Headers");
+
     public AMQPPublisherGui(){
         init();
     }
@@ -74,6 +78,7 @@ public void configure(TestElement element) {
         replyToQueue.setText(sampler.getReplyToQueue());
         correlationId.setText(sampler.getCorrelationId());
         message.setText(sampler.getMessage());
+        configureHeaders(sampler);
     }
 
     /**
@@ -105,6 +110,7 @@ public void modifyTestElement(TestElement te) {
         sampler.setMessageType(messageType.getText());
         sampler.setReplyToQueue(replyToQueue.getText());
         sampler.setCorrelationId(correlationId.getText());
+        sampler.setHeaders((Arguments) headers.createTestElement());
     }
 
     @Override
@@ -132,6 +138,7 @@ protected final void init() {
         mainPanel.add(messageType);
         mainPanel.add(replyToQueue);
         mainPanel.add(correlationId);
+        mainPanel.add(headers);
         mainPanel.add(message);
     }
 
@@ -147,6 +154,17 @@ public void clearGui() {
         messageType.setText("");
         replyToQueue.setText("");
         correlationId.setText("");
+        headers.clearGui();
         message.setText("");
     }
+
+    private void configureHeaders(AMQPPublisher sampler)
+    {
+        Arguments sampleHeaders = sampler.getHeaders();
+        if (sampleHeaders != null) {
+            headers.configure(sampleHeaders);
+        } else {
+            headers.clearGui();
+        }
+    }
 }