From e87cb7905e07372d815907239501ecee639db4cd Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Radu=20Lucu=C8=9B?= <git@radulucut.com>
Date: Fri, 9 Aug 2024 17:58:20 +0300
Subject: [PATCH] misc: allow user agent to be passed to the globalping client
 (#127)

---
 cmd/common.go                 |  5 +++++
 cmd/common_test.go            |  6 ++++++
 cmd/root.go                   |  5 +++--
 globalping/client.go          | 13 ++++++++-----
 globalping/globalping.go      |  9 ++-------
 globalping/globalping_test.go |  6 ------
 6 files changed, 24 insertions(+), 20 deletions(-)

diff --git a/cmd/common.go b/cmd/common.go
index e9a214d..4784a3a 100644
--- a/cmd/common.go
+++ b/cmd/common.go
@@ -17,6 +17,7 @@ import (
 
 	"github.com/icza/backscanner"
 	"github.com/jsdelivr/globalping-cli/globalping"
+	"github.com/jsdelivr/globalping-cli/version"
 	"github.com/shirou/gopsutil/process"
 )
 
@@ -321,3 +322,7 @@ func silenceUsageOnCreateMeasurementError(err error) bool {
 	}
 	return true
 }
+
+func getUserAgent() string {
+	return fmt.Sprintf("globalping-cli/v%s (https://github.com/jsdelivr/globalping-cli)", version.Version)
+}
diff --git a/cmd/common_test.go b/cmd/common_test.go
index 1f2ca64..8675e67 100644
--- a/cmd/common_test.go
+++ b/cmd/common_test.go
@@ -4,6 +4,7 @@ import (
 	"os"
 	"testing"
 
+	"github.com/jsdelivr/globalping-cli/version"
 	"github.com/jsdelivr/globalping-cli/view"
 	"github.com/stretchr/testify/assert"
 )
@@ -200,3 +201,8 @@ func Test_FindAndRemoveResolver_ResolverOnly(t *testing.T) {
 	assert.Equal(t, "1.1.1.1", resolver)
 	assert.Equal(t, []string{"example.com"}, argsWithoutResolver)
 }
+
+func TestUserAgent(t *testing.T) {
+	version.Version = "x.y.z"
+	assert.Equal(t, "globalping-cli/vx.y.z (https://github.com/jsdelivr/globalping-cli)", getUserAgent())
+}
diff --git a/cmd/root.go b/cmd/root.go
index d50aba3..e63b798 100644
--- a/cmd/root.go
+++ b/cmd/root.go
@@ -42,8 +42,9 @@ func Execute() {
 	}
 	t := time.NewTicker(10 * time.Second)
 	globalpingClient := globalping.NewClientWithCacheCleanup(globalping.Config{
-		APIURL:   config.GlobalpingAPIURL,
-		APIToken: config.GlobalpingToken,
+		APIURL:    config.GlobalpingAPIURL,
+		APIToken:  config.GlobalpingToken,
+		UserAgent: getUserAgent(),
 	}, t, 30)
 	globalpingProbe := probe.NewProbe()
 	viewer := view.NewViewer(ctx, printer, utime, globalpingClient)
diff --git a/globalping/client.go b/globalping/client.go
index ede2803..4468a72 100644
--- a/globalping/client.go
+++ b/globalping/client.go
@@ -22,8 +22,9 @@ type Client interface {
 }
 
 type Config struct {
-	APIURL   string
-	APIToken string
+	APIURL    string
+	APIToken  string
+	UserAgent string
 }
 
 type CacheEntry struct {
@@ -40,6 +41,7 @@ type client struct {
 	apiURL                        string
 	apiToken                      string
 	apiResponseCacheExpireSeconds int64
+	userAgent                     string
 }
 
 // NewClient creates a new client with the given configuration.
@@ -50,9 +52,10 @@ func NewClient(config Config) Client {
 		http: &http.Client{
 			Timeout: 30 * time.Second,
 		},
-		apiURL:   config.APIURL,
-		apiToken: config.APIToken,
-		cache:    map[string]*CacheEntry{},
+		apiURL:    config.APIURL,
+		apiToken:  config.APIToken,
+		userAgent: config.UserAgent,
+		cache:     map[string]*CacheEntry{},
 	}
 }
 
diff --git a/globalping/globalping.go b/globalping/globalping.go
index c13bccb..7fed187 100644
--- a/globalping/globalping.go
+++ b/globalping/globalping.go
@@ -10,7 +10,6 @@ import (
 
 	"github.com/andybalholm/brotli"
 	"github.com/jsdelivr/globalping-cli/utils"
-	"github.com/jsdelivr/globalping-cli/version"
 )
 
 var (
@@ -30,7 +29,7 @@ func (c *client) CreateMeasurement(measurement *MeasurementCreate) (*Measurement
 	if err != nil {
 		return nil, &MeasurementError{Message: "failed to create request - please report this bug"}
 	}
-	req.Header.Set("User-Agent", userAgent())
+	req.Header.Set("User-Agent", c.userAgent)
 	req.Header.Set("Accept-Encoding", "br")
 	req.Header.Set("Content-Type", "application/json")
 
@@ -146,7 +145,7 @@ func (c *client) GetMeasurementRaw(id string) ([]byte, error) {
 		return nil, &MeasurementError{Message: "failed to create request"}
 	}
 
-	req.Header.Set("User-Agent", userAgent())
+	req.Header.Set("User-Agent", c.userAgent)
 	req.Header.Set("Accept-Encoding", "br")
 
 	etag := c.getETag(id)
@@ -292,7 +291,3 @@ func DecodeHTTPTLS(tls json.RawMessage) (*HTTPTLSCertificate, error) {
 	}
 	return t, nil
 }
-
-func userAgent() string {
-	return fmt.Sprintf("globalping-cli/v%s (https://github.com/jsdelivr/globalping-cli)", version.Version)
-}
diff --git a/globalping/globalping_test.go b/globalping/globalping_test.go
index 157fa44..152694e 100644
--- a/globalping/globalping_test.go
+++ b/globalping/globalping_test.go
@@ -10,7 +10,6 @@ import (
 	"testing"
 
 	"github.com/andybalholm/brotli"
-	"github.com/jsdelivr/globalping-cli/version"
 
 	"github.com/stretchr/testify/assert"
 )
@@ -859,11 +858,6 @@ func TestFetchWithBrotli(t *testing.T) {
 	assert.Equal(t, id, m.ID)
 }
 
-func TestUserAgent(t *testing.T) {
-	version.Version = "x.y.z"
-	assert.Equal(t, "globalping-cli/vx.y.z (https://github.com/jsdelivr/globalping-cli)", userAgent())
-}
-
 // Generate server for testing
 func generateServer(json string, statusCode int) *httptest.Server {
 	server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {