mcpgrafana_test.go•21.3 kB
//go:build unit
// +build unit
package mcpgrafana
import (
"context"
"net/http"
"testing"
"github.com/go-openapi/runtime/client"
grafana_client "github.com/grafana/grafana-openapi-client-go/client"
"github.com/mark3labs/mcp-go/mcp"
"github.com/stretchr/testify/assert"
"github.com/stretchr/testify/require"
"go.opentelemetry.io/otel"
"go.opentelemetry.io/otel/attribute"
"go.opentelemetry.io/otel/codes"
sdktrace "go.opentelemetry.io/otel/sdk/trace"
"go.opentelemetry.io/otel/sdk/trace/tracetest"
)
func TestExtractIncidentClientFromEnv(t *testing.T) {
t.Setenv("GRAFANA_URL", "http://my-test-url.grafana.com/")
ctx := ExtractIncidentClientFromEnv(context.Background())
client := IncidentClientFromContext(ctx)
require.NotNil(t, client)
assert.Equal(t, "http://my-test-url.grafana.com/api/plugins/grafana-irm-app/resources/api/v1/", client.RemoteHost)
}
func TestExtractIncidentClientFromHeaders(t *testing.T) {
t.Run("no headers, no env", func(t *testing.T) {
req, err := http.NewRequest("GET", "http://example.com", nil)
require.NoError(t, err)
ctx := ExtractIncidentClientFromHeaders(context.Background(), req)
client := IncidentClientFromContext(ctx)
require.NotNil(t, client)
assert.Equal(t, "http://localhost:3000/api/plugins/grafana-irm-app/resources/api/v1/", client.RemoteHost)
})
t.Run("no headers, with env", func(t *testing.T) {
t.Setenv("GRAFANA_URL", "http://my-test-url.grafana.com/")
req, err := http.NewRequest("GET", "http://example.com", nil)
require.NoError(t, err)
ctx := ExtractIncidentClientFromHeaders(context.Background(), req)
client := IncidentClientFromContext(ctx)
require.NotNil(t, client)
assert.Equal(t, "http://my-test-url.grafana.com/api/plugins/grafana-irm-app/resources/api/v1/", client.RemoteHost)
})
t.Run("with headers, no env", func(t *testing.T) {
req, err := http.NewRequest("GET", "http://example.com", nil)
req.Header.Set(grafanaURLHeader, "http://my-test-url.grafana.com")
require.NoError(t, err)
ctx := ExtractIncidentClientFromHeaders(context.Background(), req)
client := IncidentClientFromContext(ctx)
require.NotNil(t, client)
assert.Equal(t, "http://my-test-url.grafana.com/api/plugins/grafana-irm-app/resources/api/v1/", client.RemoteHost)
})
t.Run("with headers, with env", func(t *testing.T) {
t.Setenv("GRAFANA_URL", "will-not-be-used")
req, err := http.NewRequest("GET", "http://example.com", nil)
req.Header.Set(grafanaURLHeader, "http://my-test-url.grafana.com")
require.NoError(t, err)
ctx := ExtractIncidentClientFromHeaders(context.Background(), req)
client := IncidentClientFromContext(ctx)
require.NotNil(t, client)
assert.Equal(t, "http://my-test-url.grafana.com/api/plugins/grafana-irm-app/resources/api/v1/", client.RemoteHost)
})
}
func TestExtractGrafanaInfoFromHeaders(t *testing.T) {
t.Run("no headers, no env", func(t *testing.T) {
// Explicitly clear environment variables to ensure test isolation
t.Setenv("GRAFANA_URL", "")
t.Setenv("GRAFANA_API_KEY", "")
t.Setenv("GRAFANA_SERVICE_ACCOUNT_TOKEN", "")
req, err := http.NewRequest("GET", "http://example.com", nil)
require.NoError(t, err)
ctx := ExtractGrafanaInfoFromHeaders(context.Background(), req)
config := GrafanaConfigFromContext(ctx)
assert.Equal(t, defaultGrafanaURL, config.URL)
assert.Equal(t, "", config.APIKey)
assert.Nil(t, config.BasicAuth)
})
t.Run("no headers, with env", func(t *testing.T) {
t.Setenv("GRAFANA_URL", "http://my-test-url.grafana.com")
t.Setenv("GRAFANA_API_KEY", "my-test-api-key")
req, err := http.NewRequest("GET", "http://example.com", nil)
require.NoError(t, err)
ctx := ExtractGrafanaInfoFromHeaders(context.Background(), req)
config := GrafanaConfigFromContext(ctx)
assert.Equal(t, "http://my-test-url.grafana.com", config.URL)
assert.Equal(t, "my-test-api-key", config.APIKey)
})
t.Run("no headers, with service account token", func(t *testing.T) {
t.Setenv("GRAFANA_URL", "http://my-test-url.grafana.com")
t.Setenv("GRAFANA_SERVICE_ACCOUNT_TOKEN", "my-service-account-token")
req, err := http.NewRequest("GET", "http://example.com", nil)
require.NoError(t, err)
ctx := ExtractGrafanaInfoFromHeaders(context.Background(), req)
config := GrafanaConfigFromContext(ctx)
assert.Equal(t, "http://my-test-url.grafana.com", config.URL)
assert.Equal(t, "my-service-account-token", config.APIKey)
})
t.Run("no headers, service account token takes precedence over api key", func(t *testing.T) {
t.Setenv("GRAFANA_URL", "http://my-test-url.grafana.com")
t.Setenv("GRAFANA_API_KEY", "my-deprecated-api-key")
t.Setenv("GRAFANA_SERVICE_ACCOUNT_TOKEN", "my-service-account-token")
req, err := http.NewRequest("GET", "http://example.com", nil)
require.NoError(t, err)
ctx := ExtractGrafanaInfoFromHeaders(context.Background(), req)
config := GrafanaConfigFromContext(ctx)
assert.Equal(t, "http://my-test-url.grafana.com", config.URL)
assert.Equal(t, "my-service-account-token", config.APIKey)
})
t.Run("with headers, no env", func(t *testing.T) {
req, err := http.NewRequest("GET", "http://example.com", nil)
require.NoError(t, err)
req.Header.Set(grafanaURLHeader, "http://my-test-url.grafana.com")
req.Header.Set(grafanaAPIKeyHeader, "my-test-api-key")
ctx := ExtractGrafanaInfoFromHeaders(context.Background(), req)
config := GrafanaConfigFromContext(ctx)
assert.Equal(t, "http://my-test-url.grafana.com", config.URL)
assert.Equal(t, "my-test-api-key", config.APIKey)
})
t.Run("with headers, with env", func(t *testing.T) {
// Env vars should be ignored if headers are present.
t.Setenv("GRAFANA_URL", "will-not-be-used")
t.Setenv("GRAFANA_API_KEY", "will-not-be-used")
t.Setenv("GRAFANA_SERVICE_ACCOUNT_TOKEN", "will-not-be-used")
req, err := http.NewRequest("GET", "http://example.com", nil)
require.NoError(t, err)
req.Header.Set(grafanaURLHeader, "http://my-test-url.grafana.com")
req.Header.Set(grafanaAPIKeyHeader, "my-test-api-key")
ctx := ExtractGrafanaInfoFromHeaders(context.Background(), req)
config := GrafanaConfigFromContext(ctx)
assert.Equal(t, "http://my-test-url.grafana.com", config.URL)
assert.Equal(t, "my-test-api-key", config.APIKey)
})
t.Run("no headers, with env", func(t *testing.T) {
t.Setenv("GRAFANA_USERNAME", "foo")
t.Setenv("GRAFANA_PASSWORD", "bar")
req, err := http.NewRequest("GET", "http://example.com", nil)
require.NoError(t, err)
ctx := ExtractGrafanaInfoFromHeaders(context.Background(), req)
config := GrafanaConfigFromContext(ctx)
assert.Equal(t, "foo", config.BasicAuth.Username())
password, _ := config.BasicAuth.Password()
assert.Equal(t, "bar", password)
})
t.Run("user auth with headers, no env", func(t *testing.T) {
req, err := http.NewRequest("GET", "http://example.com", nil)
req.SetBasicAuth("foo", "bar")
require.NoError(t, err)
ctx := ExtractGrafanaInfoFromHeaders(context.Background(), req)
config := GrafanaConfigFromContext(ctx)
assert.Equal(t, "foo", config.BasicAuth.Username())
password, _ := config.BasicAuth.Password()
assert.Equal(t, "bar", password)
})
t.Run("user auth with headers, with env", func(t *testing.T) {
t.Setenv("GRAFANA_USERNAME", "will-not-be-used")
t.Setenv("GRAFANA_PASSWORD", "will-not-be-used")
req, err := http.NewRequest("GET", "http://example.com", nil)
req.SetBasicAuth("foo", "bar")
require.NoError(t, err)
ctx := ExtractGrafanaInfoFromHeaders(context.Background(), req)
config := GrafanaConfigFromContext(ctx)
assert.Equal(t, "foo", config.BasicAuth.Username())
password, _ := config.BasicAuth.Password()
assert.Equal(t, "bar", password)
})
}
func TestExtractGrafanaClientPath(t *testing.T) {
t.Run("no custom path", func(t *testing.T) {
t.Setenv("GRAFANA_URL", "http://my-test-url.grafana.com/")
ctx := ExtractGrafanaClientFromEnv(context.Background())
c := GrafanaClientFromContext(ctx)
require.NotNil(t, c)
rt := c.Transport.(*client.Runtime)
assert.Equal(t, "/api", rt.BasePath)
})
t.Run("custom path", func(t *testing.T) {
t.Setenv("GRAFANA_URL", "http://my-test-url.grafana.com/grafana")
ctx := ExtractGrafanaClientFromEnv(context.Background())
c := GrafanaClientFromContext(ctx)
require.NotNil(t, c)
rt := c.Transport.(*client.Runtime)
assert.Equal(t, "/grafana/api", rt.BasePath)
})
t.Run("custom path, trailing slash", func(t *testing.T) {
t.Setenv("GRAFANA_URL", "http://my-test-url.grafana.com/grafana/")
ctx := ExtractGrafanaClientFromEnv(context.Background())
c := GrafanaClientFromContext(ctx)
require.NotNil(t, c)
rt := c.Transport.(*client.Runtime)
assert.Equal(t, "/grafana/api", rt.BasePath)
})
}
// minURL is a helper struct representing what we can extract from a constructed
// Grafana client.
type minURL struct {
host, basePath string
}
// minURLFromClient extracts some minimal amount of URL info from a Grafana client.
func minURLFromClient(c *grafana_client.GrafanaHTTPAPI) minURL {
rt := c.Transport.(*client.Runtime)
return minURL{rt.Host, rt.BasePath}
}
func TestExtractGrafanaClientFromHeaders(t *testing.T) {
t.Run("no headers, no env", func(t *testing.T) {
req, err := http.NewRequest("GET", "http://example.com", nil)
require.NoError(t, err)
ctx := ExtractGrafanaClientFromHeaders(context.Background(), req)
c := GrafanaClientFromContext(ctx)
url := minURLFromClient(c)
assert.Equal(t, "localhost:3000", url.host)
assert.Equal(t, "/api", url.basePath)
})
t.Run("no headers, with env", func(t *testing.T) {
t.Setenv("GRAFANA_URL", "http://my-test-url.grafana.com")
req, err := http.NewRequest("GET", "http://example.com", nil)
require.NoError(t, err)
ctx := ExtractGrafanaClientFromHeaders(context.Background(), req)
c := GrafanaClientFromContext(ctx)
url := minURLFromClient(c)
assert.Equal(t, "my-test-url.grafana.com", url.host)
assert.Equal(t, "/api", url.basePath)
})
t.Run("with headers, no env", func(t *testing.T) {
req, err := http.NewRequest("GET", "http://example.com", nil)
require.NoError(t, err)
req.Header.Set(grafanaURLHeader, "http://my-test-url.grafana.com")
ctx := ExtractGrafanaClientFromHeaders(context.Background(), req)
c := GrafanaClientFromContext(ctx)
url := minURLFromClient(c)
assert.Equal(t, "my-test-url.grafana.com", url.host)
assert.Equal(t, "/api", url.basePath)
})
t.Run("with headers, with env", func(t *testing.T) {
// Env vars should be ignored if headers are present.
t.Setenv("GRAFANA_URL", "will-not-be-used")
req, err := http.NewRequest("GET", "http://example.com", nil)
require.NoError(t, err)
req.Header.Set(grafanaURLHeader, "http://my-test-url.grafana.com")
ctx := ExtractGrafanaClientFromHeaders(context.Background(), req)
c := GrafanaClientFromContext(ctx)
url := minURLFromClient(c)
assert.Equal(t, "my-test-url.grafana.com", url.host)
assert.Equal(t, "/api", url.basePath)
})
}
func TestToolTracingInstrumentation(t *testing.T) {
// Set up in-memory span recorder
spanRecorder := tracetest.NewSpanRecorder()
tracerProvider := sdktrace.NewTracerProvider(
sdktrace.WithSpanProcessor(spanRecorder),
)
originalProvider := otel.GetTracerProvider()
otel.SetTracerProvider(tracerProvider)
defer otel.SetTracerProvider(originalProvider) // Restore original provider
t.Run("successful tool execution creates span with correct attributes", func(t *testing.T) {
// Clear any previous spans
spanRecorder.Reset()
// Define a simple test tool
type TestParams struct {
Message string `json:"message" jsonschema:"description=Test message"`
}
testHandler := func(ctx context.Context, args TestParams) (string, error) {
return "Hello " + args.Message, nil
}
// Create tool using MustTool (this applies our instrumentation)
tool := MustTool("test_tool", "A test tool for tracing", testHandler)
// Create context with argument logging enabled
config := GrafanaConfig{
IncludeArgumentsInSpans: true,
}
ctx := WithGrafanaConfig(context.Background(), config)
// Create a mock MCP request
request := mcp.CallToolRequest{
Params: struct {
Name string `json:"name"`
Arguments any `json:"arguments,omitempty"`
Meta *mcp.Meta `json:"_meta,omitempty"`
}{
Name: "test_tool",
Arguments: map[string]interface{}{
"message": "world",
},
},
}
// Execute the tool
result, err := tool.Handler(ctx, request)
require.NoError(t, err)
require.NotNil(t, result)
// Verify span was created
spans := spanRecorder.Ended()
require.Len(t, spans, 1)
span := spans[0]
assert.Equal(t, "mcp.tool.test_tool", span.Name())
assert.Equal(t, codes.Ok, span.Status().Code)
// Check attributes
attributes := span.Attributes()
assertHasAttribute(t, attributes, "mcp.tool.name", "test_tool")
assertHasAttribute(t, attributes, "mcp.tool.description", "A test tool for tracing")
assertHasAttribute(t, attributes, "mcp.tool.arguments", `{"message":"world"}`)
})
t.Run("tool execution error records error on span", func(t *testing.T) {
// Clear any previous spans
spanRecorder.Reset()
// Define a test tool that returns an error
type TestParams struct {
ShouldFail bool `json:"shouldFail" jsonschema:"description=Whether to fail"`
}
testHandler := func(ctx context.Context, args TestParams) (string, error) {
if args.ShouldFail {
return "", assert.AnError
}
return "success", nil
}
// Create tool
tool := MustTool("failing_tool", "A tool that can fail", testHandler)
// Create context (spans always created)
config := GrafanaConfig{}
ctx := WithGrafanaConfig(context.Background(), config)
// Create a mock MCP request that will cause failure
request := mcp.CallToolRequest{
Params: struct {
Name string `json:"name"`
Arguments any `json:"arguments,omitempty"`
Meta *mcp.Meta `json:"_meta,omitempty"`
}{
Name: "failing_tool",
Arguments: map[string]interface{}{
"shouldFail": true,
},
},
}
// Execute the tool (should fail)
result, err := tool.Handler(ctx, request)
assert.Error(t, err)
assert.Nil(t, result)
// Verify span was created and marked as error
spans := spanRecorder.Ended()
require.Len(t, spans, 1)
span := spans[0]
assert.Equal(t, "mcp.tool.failing_tool", span.Name())
assert.Equal(t, codes.Error, span.Status().Code)
assert.Equal(t, assert.AnError.Error(), span.Status().Description)
// Verify error was recorded (check events for error record)
events := span.Events()
hasErrorEvent := false
for _, event := range events {
if event.Name == "exception" {
hasErrorEvent = true
break
}
}
assert.True(t, hasErrorEvent, "Expected error event to be recorded on span")
})
t.Run("spans always created for context propagation", func(t *testing.T) {
// Clear any previous spans
spanRecorder.Reset()
// Define a simple test tool
type TestParams struct {
Message string `json:"message" jsonschema:"description=Test message"`
}
testHandler := func(ctx context.Context, args TestParams) (string, error) {
return "processed", nil
}
// Create tool
tool := MustTool("context_prop_tool", "A tool for context propagation", testHandler)
// Create context with default config (no special flags)
config := GrafanaConfig{}
ctx := WithGrafanaConfig(context.Background(), config)
// Create a mock MCP request
request := mcp.CallToolRequest{
Params: struct {
Name string `json:"name"`
Arguments any `json:"arguments,omitempty"`
Meta *mcp.Meta `json:"_meta,omitempty"`
}{
Name: "context_prop_tool",
Arguments: map[string]interface{}{
"message": "test",
},
},
}
// Execute the tool (should always create spans for context propagation)
result, err := tool.Handler(ctx, request)
require.NoError(t, err)
require.NotNil(t, result)
// Verify spans ARE always created
spans := spanRecorder.Ended()
require.Len(t, spans, 1)
span := spans[0]
assert.Equal(t, "mcp.tool.context_prop_tool", span.Name())
assert.Equal(t, codes.Ok, span.Status().Code)
})
t.Run("arguments not logged by default (PII safety)", func(t *testing.T) {
// Clear any previous spans
spanRecorder.Reset()
// Define a simple test tool
type TestParams struct {
SensitiveData string `json:"sensitiveData" jsonschema:"description=Potentially sensitive data"`
}
testHandler := func(ctx context.Context, args TestParams) (string, error) {
return "processed", nil
}
// Create tool
tool := MustTool("sensitive_tool", "A tool with sensitive data", testHandler)
// Create context with argument logging disabled (default)
config := GrafanaConfig{
IncludeArgumentsInSpans: false, // Default: safe
}
ctx := WithGrafanaConfig(context.Background(), config)
// Create a mock MCP request with potentially sensitive data
request := mcp.CallToolRequest{
Params: struct {
Name string `json:"name"`
Arguments any `json:"arguments,omitempty"`
Meta *mcp.Meta `json:"_meta,omitempty"`
}{
Name: "sensitive_tool",
Arguments: map[string]interface{}{
"sensitiveData": "user@example.com",
},
},
}
// Execute the tool (arguments should NOT be logged by default)
result, err := tool.Handler(ctx, request)
require.NoError(t, err)
require.NotNil(t, result)
// Verify span was created
spans := spanRecorder.Ended()
require.Len(t, spans, 1)
span := spans[0]
assert.Equal(t, "mcp.tool.sensitive_tool", span.Name())
assert.Equal(t, codes.Ok, span.Status().Code)
// Check that arguments are NOT logged (PII safety)
attributes := span.Attributes()
assertHasAttribute(t, attributes, "mcp.tool.name", "sensitive_tool")
assertHasAttribute(t, attributes, "mcp.tool.description", "A tool with sensitive data")
// Verify arguments are NOT present
for _, attr := range attributes {
assert.NotEqual(t, "mcp.tool.arguments", string(attr.Key), "Arguments should not be logged by default for PII safety")
}
})
t.Run("arguments logged when argument logging enabled", func(t *testing.T) {
// Clear any previous spans
spanRecorder.Reset()
// Define a simple test tool
type TestParams struct {
SafeData string `json:"safeData" jsonschema:"description=Non-sensitive data"`
}
testHandler := func(ctx context.Context, args TestParams) (string, error) {
return "processed", nil
}
// Create tool
tool := MustTool("debug_tool", "A tool for debugging", testHandler)
// Create context with argument logging enabled
config := GrafanaConfig{
IncludeArgumentsInSpans: true,
}
ctx := WithGrafanaConfig(context.Background(), config)
// Create a mock MCP request
request := mcp.CallToolRequest{
Params: struct {
Name string `json:"name"`
Arguments any `json:"arguments,omitempty"`
Meta *mcp.Meta `json:"_meta,omitempty"`
}{
Name: "debug_tool",
Arguments: map[string]interface{}{
"safeData": "debug-value",
},
},
}
// Execute the tool (arguments SHOULD be logged when flag enabled)
result, err := tool.Handler(ctx, request)
require.NoError(t, err)
require.NotNil(t, result)
// Verify span was created
spans := spanRecorder.Ended()
require.Len(t, spans, 1)
span := spans[0]
assert.Equal(t, "mcp.tool.debug_tool", span.Name())
assert.Equal(t, codes.Ok, span.Status().Code)
// Check that arguments ARE logged when flag enabled
attributes := span.Attributes()
assertHasAttribute(t, attributes, "mcp.tool.name", "debug_tool")
assertHasAttribute(t, attributes, "mcp.tool.description", "A tool for debugging")
assertHasAttribute(t, attributes, "mcp.tool.arguments", `{"safeData":"debug-value"}`)
})
}
func TestHTTPTracingConfiguration(t *testing.T) {
t.Run("HTTP tracing always enabled for context propagation", func(t *testing.T) {
// Create context (HTTP tracing always enabled)
config := GrafanaConfig{}
ctx := WithGrafanaConfig(context.Background(), config)
// Create Grafana client
client := NewGrafanaClient(ctx, "http://localhost:3000", "test-api-key", nil)
require.NotNil(t, client)
// Verify the client was created successfully (should not panic)
assert.NotNil(t, client.Transport)
})
t.Run("tracing works gracefully without OpenTelemetry configured", func(t *testing.T) {
// No OpenTelemetry tracer provider configured
// Create context (tracing always enabled for context propagation)
config := GrafanaConfig{}
ctx := WithGrafanaConfig(context.Background(), config)
// Create Grafana client (should not panic even without OTEL configured)
client := NewGrafanaClient(ctx, "http://localhost:3000", "test-api-key", nil)
require.NotNil(t, client)
// Verify the client was created successfully
assert.NotNil(t, client.Transport)
})
}
// Helper function to check if an attribute exists with expected value
func assertHasAttribute(t *testing.T, attributes []attribute.KeyValue, key string, expectedValue string) {
for _, attr := range attributes {
if string(attr.Key) == key {
assert.Equal(t, expectedValue, attr.Value.AsString())
return
}
}
t.Errorf("Expected attribute %s with value %s not found", key, expectedValue)
}