From 7a437d89ec26a714d7b653a06e8cc53a7b68c24e Mon Sep 17 00:00:00 2001
From: rogerogers <rogers@rogerogers.com>
Date: Fri, 7 Oct 2022 17:41:15 +0800
Subject: [PATCH] add test for with custom user agent of exporter/otlp

Signed-off-by: rogerogers <rogers@rogerogers.com>
---
 .../otlp/otlpmetric/otlpmetricgrpc/client_test.go | 15 +++++++++++++++
 .../otlp/otlpmetric/otlpmetrichttp/client_test.go | 15 +++++++++++++++
 .../otlp/otlptrace/otlptracegrpc/client_test.go   | 15 +++++++++++++++
 .../otlp/otlptrace/otlptracehttp/client_test.go   | 13 +++++++++++++
 4 files changed, 58 insertions(+)

diff --git a/exporters/otlp/otlpmetric/otlpmetricgrpc/client_test.go b/exporters/otlp/otlpmetric/otlpmetricgrpc/client_test.go
index 7d949a16c4e..64d3b216a2d 100644
--- a/exporters/otlp/otlpmetric/otlpmetricgrpc/client_test.go
+++ b/exporters/otlp/otlpmetric/otlpmetricgrpc/client_test.go
@@ -22,6 +22,7 @@ import (
 	"github.com/stretchr/testify/assert"
 	"github.com/stretchr/testify/require"
 	"google.golang.org/genproto/googleapis/rpc/errdetails"
+	"google.golang.org/grpc"
 	"google.golang.org/grpc/codes"
 	"google.golang.org/grpc/status"
 	"google.golang.org/protobuf/types/known/durationpb"
@@ -189,4 +190,18 @@ func TestConfig(t *testing.T) {
 		err := exp.Export(ctx, metricdata.ResourceMetrics{})
 		assert.ErrorContains(t, err, context.DeadlineExceeded.Error())
 	})
+
+	t.Run("WithCustomUserAgent", func(t *testing.T) {
+		key := "user-agent"
+		customerUserAgent := "custom-user-agent"
+		exp, coll := factoryFunc(nil, WithDialOption(grpc.WithUserAgent(customerUserAgent)))
+		t.Cleanup(coll.Shutdown)
+		ctx := context.Background()
+		require.NoError(t, exp.Export(ctx, metricdata.ResourceMetrics{}))
+		// Ensure everything is flushed.
+		require.NoError(t, exp.Shutdown(ctx))
+
+		got := coll.Headers()
+		assert.Contains(t, got[key][0], customerUserAgent)
+	})
 }
diff --git a/exporters/otlp/otlpmetric/otlpmetrichttp/client_test.go b/exporters/otlp/otlpmetric/otlpmetrichttp/client_test.go
index bc78656bbda..09a6c15d82c 100644
--- a/exporters/otlp/otlpmetric/otlpmetrichttp/client_test.go
+++ b/exporters/otlp/otlpmetric/otlpmetrichttp/client_test.go
@@ -162,4 +162,19 @@ func TestConfig(t *testing.T) {
 		assert.NoError(t, exp.Export(ctx, metricdata.ResourceMetrics{}))
 		assert.Len(t, coll.Collect().Dump(), 1)
 	})
+
+	t.Run("WithCustomUserAgent", func(t *testing.T) {
+		key := http.CanonicalHeaderKey("user-agent")
+		headers := map[string]string{key: "custom-user-agent"}
+		exp, coll := factoryFunc("", nil, WithHeaders(headers))
+		ctx := context.Background()
+		t.Cleanup(func() { require.NoError(t, coll.Shutdown(ctx)) })
+		require.NoError(t, exp.Export(ctx, metricdata.ResourceMetrics{}))
+		// Ensure everything is flushed.
+		require.NoError(t, exp.Shutdown(ctx))
+
+		got := coll.Headers()
+		require.Contains(t, got, key)
+		assert.Equal(t, got[key], []string{headers[key]})
+	})
 }
diff --git a/exporters/otlp/otlptrace/otlptracegrpc/client_test.go b/exporters/otlp/otlptrace/otlptracegrpc/client_test.go
index 063c4a987b5..395ccc28bf4 100644
--- a/exporters/otlp/otlptrace/otlptracegrpc/client_test.go
+++ b/exporters/otlp/otlptrace/otlptracegrpc/client_test.go
@@ -412,3 +412,18 @@ func TestPartialSuccess(t *testing.T) {
 	require.Contains(t, errors[0].Error(), "partially successful")
 	require.Contains(t, errors[0].Error(), "2 spans rejected")
 }
+
+func TestCustomUserAgent(t *testing.T) {
+	customUserAgent := "custom-user-agent"
+	mc := runMockCollector(t)
+	t.Cleanup(func() { require.NoError(t, mc.stop()) })
+
+	ctx := context.Background()
+	exp := newGRPCExporter(t, ctx, mc.endpoint,
+		otlptracegrpc.WithDialOption(grpc.WithUserAgent(customUserAgent)))
+	t.Cleanup(func() { require.NoError(t, exp.Shutdown(ctx)) })
+	require.NoError(t, exp.ExportSpans(ctx, roSpans))
+
+	headers := mc.getHeaders()
+	require.Contains(t, headers.Get("user-agent")[0], customUserAgent)
+}
diff --git a/exporters/otlp/otlptrace/otlptracehttp/client_test.go b/exporters/otlp/otlptrace/otlptracehttp/client_test.go
index bf497bad4be..135b6517622 100644
--- a/exporters/otlp/otlptrace/otlptracehttp/client_test.go
+++ b/exporters/otlp/otlptrace/otlptracehttp/client_test.go
@@ -42,6 +42,10 @@ var (
 		"Otel-Go-Key-1": "somevalue",
 		"Otel-Go-Key-2": "someothervalue",
 	}
+
+	customUserAgentHeader = map[string]string{
+		"user-agent": "custome-user-agent",
+	}
 )
 
 func TestEndToEnd(t *testing.T) {
@@ -142,6 +146,15 @@ func TestEndToEnd(t *testing.T) {
 				ExpectedHeaders: testHeaders,
 			},
 		},
+		{
+			name: "with custom user agent",
+			opts: []otlptracehttp.Option{
+				otlptracehttp.WithHeaders(customUserAgentHeader),
+			},
+			mcCfg: mockCollectorConfig{
+				ExpectedHeaders: customUserAgentHeader,
+			},
+		},
 	}
 
 	for _, tc := range tests {