End messages carrying data with “Data”

diff --git a/src/main/java/net/pterodactylus/fcp/FcpMessage.java b/src/main/java/net/pterodactylus/fcp/FcpMessage.java
index 921d74d..58f28a7 100644 (file)
--- a/src/main/java/net/pterodactylus/fcp/FcpMessage.java
+++ b/src/main/java/net/pterodactylus/fcp/FcpMessage.java
@@ -170,12 +170,13 @@ public class FcpMessage implements Iterable<String> {
                for (Entry<String, String> fieldEntry : fields.entrySet()) {
                        writeLine(outputStream, fieldEntry.getKey() + "=" + fieldEntry.getValue());
                }
-               writeLine(outputStream, "EndMessage");
-               outputStream.flush();
                if (payloadInputStream != null) {
+                       writeLine(outputStream, "Data");
                        FcpUtils.copy(payloadInputStream, outputStream);
-                       outputStream.flush();
+               } else {
+                       writeLine(outputStream, "EndMessage");
                }
+               outputStream.flush();
        }
 
        //

diff --git a/src/test/java/net/pterodactylus/fcp/FcpMessageTest.java b/src/test/java/net/pterodactylus/fcp/FcpMessageTest.java
new file mode 100644 (file)
index 0000000..bc0dbd1
--- /dev/null
+++ b/src/test/java/net/pterodactylus/fcp/FcpMessageTest.java
@@ -0,0 +1,35 @@
+package net.pterodactylus.fcp;
+
+import static org.hamcrest.MatcherAssert.assertThat;
+import static org.hamcrest.Matchers.contains;
+
+import java.io.ByteArrayInputStream;
+import java.io.ByteArrayOutputStream;
+import java.util.Arrays;
+import java.util.List;
+
+import org.junit.Test;
+
+/**
+ * Unit test for {@link FcpMessage}.
+ *
+ * @author <a href="mailto:david.roden@bietr.de">David Roden</a>
+ */
+public class FcpMessageTest {
+
+       private final FcpMessage fcpMessage = new FcpMessage("TestMessage");
+
+       @Test
+       public void fcpMessageWithPayloadIsTerminatedByData() throws Exception {
+               fcpMessage.setPayloadInputStream(new ByteArrayInputStream("Test".getBytes()));
+               ByteArrayOutputStream outputStream = new ByteArrayOutputStream();
+               fcpMessage.write(outputStream);
+               List<String> lines = Arrays.asList(outputStream.toString().split("\r?\n"));
+               assertThat(lines, contains(
+                       "TestMessage",
+                       "Data",
+                       "Test"
+               ));
+       }
+
+}