From d25e820327acfb38c85c8b132de5f89c1c8b68d8 Mon Sep 17 00:00:00 2001
From: Derrick Wippler <thrawn01@gmail.com>
Date: Thu, 26 May 2022 20:20:06 -0500
Subject: [PATCH] Added tests for AddOverrideHeader()

---
 httphelpers.go   |  8 +++++++-
 messages_test.go | 35 +++++++++++++++++++++++++++++++++++
 2 files changed, 42 insertions(+), 1 deletion(-)

diff --git a/httphelpers.go b/httphelpers.go
index 5156b6ed..a0249b7b 100644
--- a/httphelpers.go
+++ b/httphelpers.go
@@ -5,7 +5,6 @@ import (
 	"context"
 	"encoding/json"
 	"fmt"
-	"github.com/pkg/errors"
 	"io"
 	"io/ioutil"
 	"mime/multipart"
@@ -15,6 +14,8 @@ import (
 	"path"
 	"regexp"
 	"strings"
+
+	"github.com/pkg/errors"
 )
 
 var validURL = regexp.MustCompile(`/v[2-4].*`)
@@ -274,6 +275,11 @@ func (r *httpRequest) NewRequest(ctx context.Context, method string, payload pay
 	}
 
 	for header, value := range r.Headers {
+		// Special case, override the Host header
+		if header == "Host" {
+			req.Host = value
+			continue
+		}
 		req.Header.Add(header, value)
 	}
 	return req, nil
diff --git a/messages_test.go b/messages_test.go
index e3786be5..9f2ad8b4 100644
--- a/messages_test.go
+++ b/messages_test.go
@@ -536,6 +536,41 @@ func TestResendStored(t *testing.T) {
 	ensure.DeepEqual(t, id, exampleID)
 }
 
+func TestAddOverrideHeader(t *testing.T) {
+	const (
+		exampleDomain  = "testDomain"
+		exampleAPIKey  = "testAPIKey"
+		toUser         = "test@test.com"
+		exampleMessage = "Queue. Thank you"
+		exampleID      = "<20111114174239.25659.5817@samples.mailgun.org>"
+	)
+	srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, req *http.Request) {
+		ensure.DeepEqual(t, req.Method, http.MethodPost)
+		ensure.DeepEqual(t, req.URL.Path, fmt.Sprintf("/v3/%s/messages", exampleDomain))
+		ensure.DeepEqual(t, req.Header.Get("CustomHeader"), "custom-value")
+		ensure.DeepEqual(t, req.Host, "example.com")
+
+		rsp := fmt.Sprintf(`{"message":"%s", "id":"%s"}`, exampleMessage, exampleID)
+		fmt.Fprint(w, rsp)
+	}))
+	defer srv.Close()
+
+	mg := NewMailgun(exampleDomain, exampleAPIKey)
+	mg.SetAPIBase(srv.URL + "/v3")
+	mg.AddOverrideHeader("Host", "example.com")
+	mg.AddOverrideHeader("CustomHeader", "custom-value")
+	ctx := context.Background()
+
+	m := mg.NewMessage(fromUser, exampleSubject, exampleText, toUser)
+	m.SetRequireTLS(true)
+	m.SetSkipVerification(true)
+
+	msg, id, err := mg.Send(ctx, m)
+	ensure.Nil(t, err)
+	ensure.DeepEqual(t, msg, exampleMessage)
+	ensure.DeepEqual(t, id, exampleID)
+}
+
 func TestSendTLSOptions(t *testing.T) {
 	const (
 		exampleDomain  = "testDomain"