From be75ef9c586400e333d6bb87ff5c4db84f0f8318 Mon Sep 17 00:00:00 2001
From: TwinProduction <twin@twinnation.org>
Date: Wed, 30 Dec 2020 19:56:12 -0500
Subject: [PATCH] Work on #58: Add default "User-Agent: Gatus/1.0" header

---
 core/service.go      | 10 ++++++++++
 core/service_test.go | 26 ++++++++++++++++++++++++++
 2 files changed, 36 insertions(+)

diff --git a/core/service.go b/core/service.go
index 0481e5ca..7880f6ff 100644
--- a/core/service.go
+++ b/core/service.go
@@ -20,6 +20,12 @@ const (
 
 	// ContentTypeHeader is the name of the header used to specify the content type
 	ContentTypeHeader = "Content-Type"
+
+	// UserAgentHeader is the name of the header used to specify the request's user agent
+	UserAgentHeader = "User-Agent"
+
+	// GatusUserAgent is the default user agent that Gatus uses to send requests.
+	GatusUserAgent = "Gatus/1.0"
 )
 
 var (
@@ -90,6 +96,10 @@ func (service *Service) ValidateAndSetDefaults() {
 	if len(service.Headers) == 0 {
 		service.Headers = make(map[string]string)
 	}
+	// Automatically add user agent header if there isn't one specified in the service configuration
+	if _, userAgentHeaderExists := service.Headers[UserAgentHeader]; !userAgentHeaderExists {
+		service.Headers[UserAgentHeader] = GatusUserAgent
+	}
 	// Automatically add "Content-Type: application/json" header if there's no Content-Type set
 	// and service.GraphQL is set to true
 	if _, contentTypeHeaderExists := service.Headers[ContentTypeHeader]; !contentTypeHeaderExists && service.GraphQL {
diff --git a/core/service_test.go b/core/service_test.go
index 23a3b76f..fa908165 100644
--- a/core/service_test.go
+++ b/core/service_test.go
@@ -130,6 +130,32 @@ func TestService_buildHTTPRequest(t *testing.T) {
 	if request.Host != "twinnation.org" {
 		t.Error("request.Host should've been twinnation.org, but was", request.Host)
 	}
+	if userAgent := request.Header.Get("User-Agent"); userAgent != GatusUserAgent {
+		t.Errorf("request.Header.Get(User-Agent) should've been %s, but was %s", GatusUserAgent, userAgent)
+	}
+}
+
+func TestService_buildHTTPRequestWithCustomUserAgent(t *testing.T) {
+	condition := Condition("[STATUS] == 200")
+	service := Service{
+		Name:       "TwiNNatioN",
+		URL:        "https://twinnation.org/health",
+		Conditions: []*Condition{&condition},
+		Headers: map[string]string{
+			"User-Agent": "Test/2.0",
+		},
+	}
+	service.ValidateAndSetDefaults()
+	request := service.buildHTTPRequest()
+	if request.Method != "GET" {
+		t.Error("request.Method should've been GET, but was", request.Method)
+	}
+	if request.Host != "twinnation.org" {
+		t.Error("request.Host should've been twinnation.org, but was", request.Host)
+	}
+	if userAgent := request.Header.Get("User-Agent"); userAgent != "Test/2.0" {
+		t.Errorf("request.Header.Get(User-Agent) should've been %s, but was %s", "Test/2.0", userAgent)
+	}
 }
 
 func TestService_buildHTTPRequestWithHostHeader(t *testing.T) {