From d2638a0457866ab41ae9e64cc9280edfbd042c77 Mon Sep 17 00:00:00 2001
From: James Kraus <jamespkraus@gmail.com>
Date: Sun, 26 Mar 2017 04:23:03 -0400
Subject: [PATCH] Bug fix - Allow adding custom headers with non-array value
 (#255)

* Bug fix - Allow adding custom headers with non-array value

* Remove whitespace

* Functional test for adding custom headers (either singular or multiple)

* Style fixes for functional test
---
 src/Mailgun/Messages/MessageBuilder.php       | 11 ++++-
 tests/Functional/MessageBuilderHeaderTest.php | 41 +++++++++++++++++++
 tests/Messages/MessageBuilderTest.php         | 11 ++++-
 3 files changed, 61 insertions(+), 2 deletions(-)
 create mode 100644 tests/Functional/MessageBuilderHeaderTest.php

diff --git a/src/Mailgun/Messages/MessageBuilder.php b/src/Mailgun/Messages/MessageBuilder.php
index a9139a5..f63f183 100644
--- a/src/Mailgun/Messages/MessageBuilder.php
+++ b/src/Mailgun/Messages/MessageBuilder.php
@@ -248,7 +248,16 @@ class MessageBuilder
         if (!preg_match('/^h:/i', $headerName)) {
             $headerName = 'h:'.$headerName;
         }
-        $this->message[$headerName] = [$headerData];
+
+        if (array_key_exists($headerName, $this->message)) {
+            if (is_array($this->message[$headerName])) {
+                $this->message[$headerName][] = $headerData;
+            } else {
+                $this->message[$headerName] = [$this->message[$headerName], $headerData];
+            }
+        } else {
+            $this->message[$headerName] = $headerData;
+        }
 
         return $this->message[$headerName];
     }
diff --git a/tests/Functional/MessageBuilderHeaderTest.php b/tests/Functional/MessageBuilderHeaderTest.php
new file mode 100644
index 0000000..0a61dd7
--- /dev/null
+++ b/tests/Functional/MessageBuilderHeaderTest.php
@@ -0,0 +1,41 @@
+<?php
+
+/*
+ * Copyright (C) 2013-2016 Mailgun
+ *
+ * This software may be modified and distributed under the terms
+ * of the MIT license. See the LICENSE file for details.
+ */
+
+namespace Mailgun\Tests\Functional;
+
+/**
+ * @author James Kraus <jkraus@imagineteam.com>
+ */
+class MessageBuilderHeaderTest extends \PHPUnit_Framework_TestCase
+{
+    public function testSimpleExample()
+    {
+        $messageValidator = function ($headers) {
+            $this->assertContains(['name' => 'h:My-Singular-Header', 'contents' => '123'], $headers);
+            $this->assertContains(['name' => 'h:My-Plural-Header[0]', 'contents' => '123'], $headers);
+            $this->assertContains(['name' => 'h:My-Plural-Header[1]', 'contents' => '456'], $headers);
+        };
+
+        // Create the mocked mailgun client.
+        $mailgun = MockedMailgun::create($this, 'POST', 'domain/messages', [], $messageValidator);
+
+        $builder = $mailgun->MessageBuilder();
+
+        $builder->addCustomHeader('My-Singular-Header', '123');
+        $builder->addCustomHeader('My-Plural-Header', '123');
+        $builder->addCustomHeader('My-Plural-Header', '456');
+
+        $builder->setFromAddress('from@example.com');
+        $builder->addToRecipient('to@example.com');
+        $builder->setSubject('Foo');
+        $builder->setTextBody('Bar');
+
+        $mailgun->sendMessage('domain', $builder->getMessage());
+    }
+}
diff --git a/tests/Messages/MessageBuilderTest.php b/tests/Messages/MessageBuilderTest.php
index 70eb179..6efadde 100644
--- a/tests/Messages/MessageBuilderTest.php
+++ b/tests/Messages/MessageBuilderTest.php
@@ -133,7 +133,16 @@ class MessageBuilderTest extends \Mailgun\Tests\MailgunTestCase
         $message = $this->client->MessageBuilder();
         $message->addCustomHeader('My-Header', '123');
         $messageObj = $message->getMessage();
-        $this->assertEquals(['h:My-Header' => ['123']], $messageObj);
+        $this->assertEquals(['h:My-Header' => '123'], $messageObj);
+    }
+
+    public function testAddMultipleCustomHeader()
+    {
+        $message = $this->client->MessageBuilder();
+        $message->addCustomHeader('My-Header', '123');
+        $message->addCustomHeader('My-Header', '456');
+        $messageObj = $message->getMessage();
+        $this->assertEquals(['h:My-Header' => ['123', '456']], $messageObj);
     }
 
     public function testSetTextBody()