diff --git a/pkg/tsdb/tempo/util/stream_util_test.go b/pkg/tsdb/tempo/util/stream_util_test.go new file mode 100644 index 00000000000..2a91433d139 --- /dev/null +++ b/pkg/tsdb/tempo/util/stream_util_test.go @@ -0,0 +1,107 @@ +package util + +import ( + "context" + "testing" + + "github.com/grafana/grafana-plugin-sdk-go/backend" + "github.com/grafana/grafana-plugin-sdk-go/backend/useragent" + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" + "google.golang.org/grpc/metadata" +) + +func TestAppendHeadersToOutgoingContext_AppendsHeadersAndUserAgent(t *testing.T) { + ctx := context.TODO() + ua, err := useragent.New("10.0.0", "linux", "amd64") + require.NoError(t, err) + ctx = backend.WithUserAgent(ctx, ua) + ctx = metadata.NewOutgoingContext(ctx, metadata.Pairs("Existing", "one")) + + req := &backend.RunStreamRequest{ + Headers: map[string]string{ + "X-Test": "value", + }, + } + + out := AppendHeadersToOutgoingContext(ctx, req) + outgoingMD, ok := metadata.FromOutgoingContext(out) + require.True(t, ok) + assert.Equal(t, []string{"value"}, outgoingMD.Get("x-test")) + assert.Equal(t, []string{ua.String()}, outgoingMD.Get("user-agent")) + assert.Equal(t, []string{"one"}, outgoingMD.Get("existing")) +} + +func TestSetHeadersFromIncomingContext_MergesTeamAndClientHeaders(t *testing.T) { + jsonData := []byte(`{ + "teamHttpHeaders": { + "headers": { + "101": [ + {"header": "X-Team", "value": "team-value"}, + {"header": "X-Shared", "value": "team-wins"} + ] + } + }, + "httpHeaderName1": "X-Client", + "httpHeaderName2": "X-Shared" + }`) + + pluginCtx := backend.PluginContext{ + DataSourceInstanceSettings: &backend.DataSourceInstanceSettings{ + JSONData: jsonData, + DecryptedSecureJSONData: map[string]string{ + "httpHeaderValue1": "client-value", + "httpHeaderValue2": "client-overridden", + }, + }, + } + + ctx := backend.WithPluginContext(context.Background(), pluginCtx) + headers, err := SetHeadersFromIncomingContext(ctx) + require.NoError(t, err) + + expected := map[string]string{ + "X-Client": "client-value", + "X-Team": "team-value", + "X-Shared": "team-wins", + } + assert.Equal(t, expected, headers) +} + +func TestGetTeamHTTPHeaders_NoTeamHeaders(t *testing.T) { + pluginCtx := backend.PluginContext{ + DataSourceInstanceSettings: &backend.DataSourceInstanceSettings{ + JSONData: []byte(`{"httpHeaderName1": "X-Client"}`), + }, + } + + headers, err := getTeamHTTPHeaders(pluginCtx) + require.NoError(t, err) + assert.Empty(t, headers) +} + +func TestGetClientOptionsHeaders_ParsesHeaders(t *testing.T) { + pluginCtx := backend.PluginContext{ + DataSourceInstanceSettings: &backend.DataSourceInstanceSettings{ + JSONData: []byte(`{"httpHeaderName1": "X-Client"}`), + DecryptedSecureJSONData: map[string]string{ + "httpHeaderValue1": "client-value", + }, + }, + } + + headers, err := getClientOptionsHeaders(context.Background(), pluginCtx) + require.NoError(t, err) + assert.Equal(t, map[string]string{"X-Client": "client-value"}, headers) +} + +func TestGetClientOptionsHeaders_InvalidJSON(t *testing.T) { + pluginCtx := backend.PluginContext{ + DataSourceInstanceSettings: &backend.DataSourceInstanceSettings{ + JSONData: []byte("{"), + }, + } + + _, err := getClientOptionsHeaders(context.Background(), pluginCtx) + require.Error(t, err) +}