Skip to content

Commit 2e9e0aa

Browse files
sd2ksd2k
committed
Add integration tags and test in CI
1 parent 60c5c8b commit 2e9e0aa

File tree

6 files changed

+126
-41
lines changed

6 files changed

+126
-41
lines changed

‎.github/workflows/test.yml

+1-1
Original file line numberDiff line numberDiff line change
@@ -39,4 +39,4 @@ jobs:
3939
cache: true
4040

4141
- name: Run tests
42-
run: go test -v ./...
42+
run: make test-all

‎Makefile

+5-1
Original file line numberDiff line numberDiff line change
@@ -10,5 +10,9 @@ lint:
1010
test:
1111
go test ./...
1212

13+
.PHONY: test-all
14+
test-all:
15+
go test -v -tags integration ./...
16+
1317
.PHONY: run
14-
go run ./...
18+
go run ./...

‎tools/datasources_test.go

+11
Original file line numberDiff line numberDiff line change
@@ -1,3 +1,7 @@
1+
// Requires a Prometheus instance running on localhost:9090.
2+
// Run with `go test -tags integration`.
3+
// go:build integration
4+
15
package tools
26

37
import (
@@ -14,6 +18,9 @@ import (
1418
"github.com/stretchr/testify/require"
1519
)
1620

21+
// newTestContext creates a new context with the Grafana URL and API key
22+
// from the environment variables GRAFANA_URL and GRAFANA_API_KEY.
23+
// TODO: move this to a shared file.
1724
func newTestContext() context.Context {
1825
cfg := client.DefaultTransportConfig()
1926
cfg.Host = "localhost:3000"
@@ -32,6 +39,10 @@ func newTestContext() context.Context {
3239
}
3340
}
3441

42+
if apiKey := os.Getenv("GRAFANA_API_KEY"); apiKey != "" {
43+
cfg.APIKey = apiKey
44+
}
45+
3546
client := client.NewHTTPClientWithConfig(strfmt.Default, cfg)
3647
return mcpgrafana.WithGrafanaClient(context.Background(), client)
3748
}

‎tools/incident.go

+2-6
Original file line numberDiff line numberDiff line change
@@ -17,7 +17,7 @@ type ListIncidentsParams struct {
1717
Status string `json:"status" jsonschema:"description=The status of the incidents to include"`
1818
}
1919

20-
func listIncidents(ctx context.Context, args ListIncidentsParams) (*mcp.CallToolResult, error) {
20+
func listIncidents(ctx context.Context, args ListIncidentsParams) (*incident.QueryIncidentsResponse, error) {
2121
c := mcpgrafana.IncidentClientFromContext(ctx)
2222
is := incident.NewIncidentsService(c)
2323
query := ""
@@ -37,11 +37,7 @@ func listIncidents(ctx context.Context, args ListIncidentsParams) (*mcp.CallTool
3737
if err != nil {
3838
return nil, fmt.Errorf("list incidents: %w", err)
3939
}
40-
b, err := json.Marshal(incidents.Incidents)
41-
if err != nil {
42-
return nil, fmt.Errorf("marshal incidents: %w", err)
43-
}
44-
return mcp.NewToolResultText(string(b)), nil
40+
return incidents, nil
4541
}
4642

4743
var ListIncidents = mcpgrafana.MustTool(

‎tools/prometheus.go

+13-33
Original file line numberDiff line numberDiff line change
@@ -2,18 +2,17 @@ package tools
22

33
import (
44
"context"
5-
"encoding/json"
65
"fmt"
76
"regexp"
87
"strings"
98
"time"
109

1110
mcpgrafana "github.com/grafana/mcp-grafana"
12-
"github.com/mark3labs/mcp-go/mcp"
1311
"github.com/mark3labs/mcp-go/server"
1412
"github.com/prometheus/client_golang/api"
1513
promv1 "github.com/prometheus/client_golang/api/prometheus/v1"
1614
"github.com/prometheus/common/config"
15+
"github.com/prometheus/common/model"
1716
)
1817

1918
func promClientFromContext(ctx context.Context, uid string) (promv1.API, error) {
@@ -75,7 +74,7 @@ type QueryPrometheusParams struct {
7574
QueryType string `json:"queryType,omitempty" jsonschema:"description=The type of query to use. Either 'range' or 'instant'"`
7675
}
7776

78-
func queryPrometheus(ctx context.Context, args QueryPrometheusParams) (*mcp.CallToolResult, error) {
77+
func queryPrometheus(ctx context.Context, args QueryPrometheusParams) (model.Value, error) {
7978
promClient, err := promClientFromContext(ctx, args.DatasourceUID)
8079
if err != nil {
8180
return nil, fmt.Errorf("getting Prometheus client: %w", err)
@@ -110,23 +109,13 @@ func queryPrometheus(ctx context.Context, args QueryPrometheusParams) (*mcp.Call
110109
if err != nil {
111110
return nil, fmt.Errorf("querying Prometheus range: %w", err)
112111
}
113-
114-
b, err := json.Marshal(result)
115-
if err != nil {
116-
return nil, fmt.Errorf("marshalling Prometheus query result: %w", err)
117-
}
118-
return mcp.NewToolResultText(string(b)), nil
112+
return result, nil
119113
} else if queryType == "instant" {
120114
result, _, err := promClient.Query(ctx, args.Expr, startTime)
121115
if err != nil {
122116
return nil, fmt.Errorf("querying Prometheus instant: %w", err)
123117
}
124-
125-
b, err := json.Marshal(result)
126-
if err != nil {
127-
return nil, fmt.Errorf("marshalling Prometheus query result: %w", err)
128-
}
129-
return mcp.NewToolResultText(string(b)), nil
118+
return result, nil
130119
}
131120

132121
return nil, fmt.Errorf("invalid query type: %s", queryType)
@@ -145,7 +134,7 @@ type ListPrometheusMetricNamesParams struct {
145134
Page int `json:"page,omitempty" jsonschema:"description=The page number to return"`
146135
}
147136

148-
func listPrometheusMetricNames(ctx context.Context, args ListPrometheusMetricNamesParams) (*mcp.CallToolResult, error) {
137+
func listPrometheusMetricNames(ctx context.Context, args ListPrometheusMetricNamesParams) ([]string, error) {
149138
promClient, err := promClientFromContext(ctx, args.DatasourceUID)
150139
if err != nil {
151140
return nil, fmt.Errorf("getting Prometheus client: %w", err)
@@ -196,11 +185,7 @@ func listPrometheusMetricNames(ctx context.Context, args ListPrometheusMetricNam
196185
matches = matches[start:end]
197186
}
198187

199-
b, err := json.Marshal(matches)
200-
if err != nil {
201-
return nil, fmt.Errorf("marshalling Prometheus metric names: %w", err)
202-
}
203-
return mcp.NewToolResultText(string(b)), nil
188+
return matches, nil
204189
}
205190

206191
var ListPrometheusMetricNames = mcpgrafana.MustTool(
@@ -223,6 +208,9 @@ func (s Selector) String() string {
223208
b := strings.Builder{}
224209
b.WriteRune('{')
225210
for i, f := range s.Filters {
211+
if f.Type == "" {
212+
f.Type = "="
213+
}
226214
b.WriteString(fmt.Sprintf(`%s%s'%s'`, f.Name, f.Type, f.Value))
227215
if i < len(s.Filters)-1 {
228216
b.WriteString(", ")
@@ -240,7 +228,7 @@ type ListPrometheusLabelNamesParams struct {
240228
Limit int `json:"limit,omitempty" jsonschema:"description=Optionally, the maximum number of results to return"`
241229
}
242230

243-
func listPrometheusLabelNames(ctx context.Context, args ListPrometheusLabelNamesParams) (*mcp.CallToolResult, error) {
231+
func listPrometheusLabelNames(ctx context.Context, args ListPrometheusLabelNamesParams) ([]string, error) {
244232
promClient, err := promClientFromContext(ctx, args.DatasourceUID)
245233
if err != nil {
246234
return nil, fmt.Errorf("getting Prometheus client: %w", err)
@@ -278,11 +266,7 @@ func listPrometheusLabelNames(ctx context.Context, args ListPrometheusLabelNames
278266
labelNames = labelNames[:limit]
279267
}
280268

281-
b, err := json.Marshal(labelNames)
282-
if err != nil {
283-
return nil, fmt.Errorf("marshalling Prometheus label names: %w", err)
284-
}
285-
return mcp.NewToolResultText(string(b)), nil
269+
return labelNames, nil
286270
}
287271

288272
var ListPrometheusLabelNames = mcpgrafana.MustTool(
@@ -300,7 +284,7 @@ type ListPrometheusLabelValuesParams struct {
300284
Limit int `json:"limit,omitempty" jsonschema:"description=Optionally, the maximum number of results to return"`
301285
}
302286

303-
func listPrometheusLabelValues(ctx context.Context, args ListPrometheusLabelValuesParams) (*mcp.CallToolResult, error) {
287+
func listPrometheusLabelValues(ctx context.Context, args ListPrometheusLabelValuesParams) (model.LabelValues, error) {
304288
promClient, err := promClientFromContext(ctx, args.DatasourceUID)
305289
if err != nil {
306290
return nil, fmt.Errorf("getting Prometheus client: %w", err)
@@ -338,11 +322,7 @@ func listPrometheusLabelValues(ctx context.Context, args ListPrometheusLabelValu
338322
labelValues = labelValues[:limit]
339323
}
340324

341-
b, err := json.Marshal(labelValues)
342-
if err != nil {
343-
return nil, fmt.Errorf("marshalling Prometheus label values: %w", err)
344-
}
345-
return mcp.NewToolResultText(string(b)), nil
325+
return labelValues, nil
346326
}
347327

348328
var ListPrometheusLabelValues = mcpgrafana.MustTool(

‎tools/prometheus_test.go

+94
Original file line numberDiff line numberDiff line change
@@ -1,8 +1,15 @@
1+
// Requires a Prometheus instance running on localhost:9090.
2+
// Run with `go test -tags integration`.
3+
//go:build integration
4+
15
package tools
26

37
import (
8+
"fmt"
49
"testing"
10+
"time"
511

12+
"github.com/prometheus/common/model"
613
"github.com/stretchr/testify/assert"
714
"github.com/stretchr/testify/require"
815
)
@@ -16,4 +23,91 @@ func TestPrometheusTools(t *testing.T) {
1623
require.NoError(t, err)
1724
assert.Len(t, result, 10)
1825
})
26+
27+
t.Run("list prometheus metric names", func(t *testing.T) {
28+
ctx := newTestContext()
29+
result, err := listPrometheusMetricNames(ctx, ListPrometheusMetricNamesParams{
30+
DatasourceUID: "prometheus",
31+
Regex: ".*",
32+
Limit: 10,
33+
})
34+
require.NoError(t, err)
35+
assert.Len(t, result, 10)
36+
})
37+
38+
t.Run("list prometheus label names", func(t *testing.T) {
39+
ctx := newTestContext()
40+
result, err := listPrometheusLabelNames(ctx, ListPrometheusLabelNamesParams{
41+
DatasourceUID: "prometheus",
42+
Matches: []Selector{
43+
{
44+
Filters: []LabelMatcher{
45+
{Name: "job", Value: "prometheus"},
46+
},
47+
},
48+
},
49+
Limit: 10,
50+
})
51+
require.NoError(t, err)
52+
assert.Len(t, result, 10)
53+
})
54+
55+
t.Run("list prometheus label values", func(t *testing.T) {
56+
ctx := newTestContext()
57+
result, err := listPrometheusLabelValues(ctx, ListPrometheusLabelValuesParams{
58+
DatasourceUID: "prometheus",
59+
LabelName: "job",
60+
Matches: []Selector{
61+
{
62+
Filters: []LabelMatcher{
63+
{Name: "job", Value: "prometheus"},
64+
},
65+
},
66+
},
67+
})
68+
require.NoError(t, err)
69+
assert.Len(t, result, 1)
70+
})
71+
}
72+
73+
func TestPrometheusQueries(t *testing.T) {
74+
t.Run("query prometheus range", func(t *testing.T) {
75+
end := time.Now()
76+
start := end.Add(-10 * time.Minute)
77+
for _, step := range []int{15, 60, 300} {
78+
t.Run(fmt.Sprintf("step=%d", step), func(t *testing.T) {
79+
ctx := newTestContext()
80+
result, err := queryPrometheus(ctx, QueryPrometheusParams{
81+
DatasourceUID: "prometheus",
82+
Expr: "up",
83+
StartRFC3339: start.Format(time.RFC3339),
84+
EndRFC3339: end.Format(time.RFC3339),
85+
StepSeconds: step,
86+
QueryType: "range",
87+
})
88+
require.NoError(t, err)
89+
matrix := result.(model.Matrix)
90+
require.Len(t, matrix, 1)
91+
expectedLen := int(end.Sub(start).Seconds()/float64(step)) + 1
92+
assert.Len(t, matrix[0].Values, expectedLen)
93+
assert.Less(t, matrix[0].Values[0].Timestamp.Sub(model.TimeFromUnix(start.Unix())), time.Duration(step)*time.Second)
94+
assert.Equal(t, matrix[0].Metric["__name__"], model.LabelValue("up"))
95+
})
96+
}
97+
})
98+
99+
t.Run("query prometheus instant", func(t *testing.T) {
100+
ctx := newTestContext()
101+
result, err := queryPrometheus(ctx, QueryPrometheusParams{
102+
DatasourceUID: "prometheus",
103+
Expr: "up",
104+
StartRFC3339: time.Now().Format(time.RFC3339),
105+
QueryType: "instant",
106+
})
107+
require.NoError(t, err)
108+
scalar := result.(model.Vector)
109+
assert.Equal(t, scalar[0].Value, model.SampleValue(1))
110+
assert.Equal(t, scalar[0].Timestamp, model.TimeFromUnix(time.Now().Unix()))
111+
assert.Equal(t, scalar[0].Metric["__name__"], model.LabelValue("up"))
112+
})
19113
}

0 commit comments

Comments
 (0)