Skip to content

Commit fc96c70

Browse files
RossTarrantRossTarrant
committed
Add unit tests
1 parent acf6dbc commit fc96c70

5 files changed

Lines changed: 245 additions & 0 deletions

File tree

‎internal/toolsnaps/toolsnaps_test.go‎

Lines changed: 5 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -163,6 +163,11 @@ func TestSortJSONKeys(t *testing.T) {
163163
input: `{"name": "test", "properties": {"repo": {"type": "string"}, "owner": {"type": "string"}, "page": {"type": "number"}}}`,
164164
expected: "{\n \"name\": \"test\",\n \"properties\": {\n \"owner\": {\n \"type\": \"string\"\n },\n \"page\": {\n \"type\": \"number\"\n },\n \"repo\": {\n \"type\": \"string\"\n }\n }\n}",
165165
},
166+
{
167+
name: "strips top-level output schema",
168+
input: `{"name": "test", "outputSchema": {"type": "object", "properties": {"id": {"type": "string"}}}, "inputSchema": {"type": "object"}}`,
169+
expected: "{\n \"inputSchema\": {\n \"type\": \"object\"\n },\n \"name\": \"test\"\n}",
170+
},
166171
}
167172

168173
for _, tt := range tests {

‎pkg/github/feature_flags_test.go‎

Lines changed: 6 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -159,6 +159,12 @@ func TestResolveFeatureFlags(t *testing.T) {
159159
insidersMode: false,
160160
expectedFlags: []string{MCPAppsFeatureFlag},
161161
},
162+
{
163+
name: "explicit output schemas feature enabled",
164+
enabledFeatures: []string{FeatureFlagOutputSchemas},
165+
insidersMode: false,
166+
expectedFlags: []string{FeatureFlagOutputSchemas},
167+
},
162168
{
163169
name: "insiders mode enables insiders flags",
164170
enabledFeatures: nil,

‎pkg/github/output_schema_test.go‎

Lines changed: 91 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,91 @@
1+
package github
2+
3+
import (
4+
"context"
5+
"encoding/json"
6+
"testing"
7+
8+
"github.com/stretchr/testify/assert"
9+
"github.com/stretchr/testify/require"
10+
)
11+
12+
type outputSchemaTestPayload struct {
13+
Message string `json:"message"`
14+
}
15+
16+
func TestMustOutputSchema(t *testing.T) {
17+
schema := MustOutputSchema[outputSchemaTestPayload]()
18+
19+
require.NotNil(t, schema)
20+
assert.Equal(t, "object", schema.Type)
21+
assert.Contains(t, schema.Properties, "message")
22+
}
23+
24+
func TestMustOutputSchemaPanicsForNonObject(t *testing.T) {
25+
require.Panics(t, func() {
26+
MustOutputSchema[string]()
27+
})
28+
}
29+
30+
func TestStructuredTextResultStructuredContentFeatureFlag(t *testing.T) {
31+
textValue := map[string]string{"message": "text"}
32+
structuredContent := outputSchemaTestPayload{Message: "structured"}
33+
34+
tests := []struct {
35+
name string
36+
featureEnabled bool
37+
wantStructured bool
38+
wantFeatureCheck bool
39+
}{
40+
{
41+
name: "omits structured content by default",
42+
wantStructured: false,
43+
wantFeatureCheck: true,
44+
},
45+
{
46+
name: "includes structured content when output schemas are enabled",
47+
featureEnabled: true,
48+
wantStructured: true,
49+
wantFeatureCheck: true,
50+
},
51+
}
52+
53+
for _, tt := range tests {
54+
t.Run(tt.name, func(t *testing.T) {
55+
var checkedFeature string
56+
deps := BaseDeps{
57+
featureChecker: func(_ context.Context, flag string) (bool, error) {
58+
checkedFeature = flag
59+
return tt.featureEnabled, nil
60+
},
61+
}
62+
63+
result, err := structuredTextResult(context.Background(), deps, textValue, structuredContent)
64+
require.NoError(t, err)
65+
require.NotNil(t, result)
66+
67+
textContent := getTextResult(t, result)
68+
var gotText map[string]string
69+
require.NoError(t, json.Unmarshal([]byte(textContent.Text), &gotText))
70+
assert.Equal(t, textValue, gotText)
71+
72+
if tt.wantStructured {
73+
assert.Equal(t, structuredContent, result.StructuredContent)
74+
} else {
75+
assert.Nil(t, result.StructuredContent)
76+
}
77+
if tt.wantFeatureCheck {
78+
assert.Equal(t, FeatureFlagOutputSchemas, checkedFeature)
79+
}
80+
})
81+
}
82+
}
83+
84+
func TestStructuredTextResultMarshalError(t *testing.T) {
85+
_, err := structuredTextResult(context.Background(), BaseDeps{}, map[string]any{
86+
"invalid": func() {},
87+
}, nil)
88+
89+
require.Error(t, err)
90+
assert.Contains(t, err.Error(), "failed to marshal response")
91+
}

‎pkg/github/tools_validation_test.go‎

Lines changed: 52 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -4,6 +4,7 @@ import (
44
"testing"
55

66
"github.com/github/github-mcp-server/pkg/inventory"
7+
"github.com/google/jsonschema-go/jsonschema"
78
"github.com/stretchr/testify/assert"
89
"github.com/stretchr/testify/require"
910
)
@@ -167,6 +168,57 @@ func TestAllToolsHaveHandlerFunc(t *testing.T) {
167168
}
168169
}
169170

171+
func TestExpectedToolsHaveOutputSchemas(t *testing.T) {
172+
expected := map[string]bool{
173+
"get_commit": true,
174+
"get_label": true,
175+
"get_latest_release": true,
176+
"get_me": true,
177+
"get_release_by_tag": true,
178+
"get_team_members": true,
179+
"get_teams": true,
180+
"list_branches": true,
181+
"list_commits": true,
182+
"list_issue_types": true,
183+
"list_issues": true,
184+
"list_pull_requests": true,
185+
"list_releases": true,
186+
"list_tags": true,
187+
"search_code": true,
188+
"search_issues": true,
189+
"search_orgs": true,
190+
"search_pull_requests": true,
191+
"search_repositories": true,
192+
"search_users": true,
193+
}
194+
seen := make(map[string]bool, len(expected))
195+
196+
for _, tool := range AllTools(stubTranslation) {
197+
t.Run(tool.Tool.Name, func(t *testing.T) {
198+
assert.Nil(t, tool.Tool.OutputSchema,
199+
"Tool %q should attach output schema only at registration time", tool.Tool.Name)
200+
201+
if tool.OutputSchema == nil {
202+
assert.False(t, expected[tool.Tool.Name],
203+
"Tool %q should have a feature-gated output schema", tool.Tool.Name)
204+
return
205+
}
206+
207+
assert.True(t, expected[tool.Tool.Name],
208+
"Tool %q has an unexpected output schema", tool.Tool.Name)
209+
schema, ok := tool.OutputSchema.(*jsonschema.Schema)
210+
require.True(t, ok, "Tool %q output schema should be *jsonschema.Schema", tool.Tool.Name)
211+
assert.Equal(t, "object", schema.Type,
212+
"Tool %q output schema should be an object schema", tool.Tool.Name)
213+
seen[tool.Tool.Name] = true
214+
})
215+
}
216+
217+
for name := range expected {
218+
assert.True(t, seen[name], "Expected output schema for tool %q", name)
219+
}
220+
}
221+
170222
// TestToolsetMetadataConsistency ensures tools in the same toolset have consistent descriptions
171223
func TestToolsetMetadataConsistency(t *testing.T) {
172224
tools := AllTools(stubTranslation)

‎pkg/inventory/registry_test.go‎

Lines changed: 91 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1200,6 +1200,97 @@ func TestFeatureFlagError(t *testing.T) {
12001200
}
12011201
}
12021202

1203+
func TestRegisterToolsOutputSchemaFeatureFlag(t *testing.T) {
1204+
outputSchema := map[string]any{
1205+
"type": "object",
1206+
"properties": map[string]any{
1207+
"id": map[string]any{"type": "string"},
1208+
},
1209+
}
1210+
tools := []ServerTool{
1211+
mockTool("tool_with_schema", "toolset1", true).WithOutputSchema(outputSchema),
1212+
}
1213+
1214+
tests := []struct {
1215+
name string
1216+
featureChecker FeatureFlagChecker
1217+
wantOutputSchema bool
1218+
}{
1219+
{
1220+
name: "omits output schema by default",
1221+
wantOutputSchema: false,
1222+
},
1223+
{
1224+
name: "includes output schema when feature is enabled",
1225+
featureChecker: func(_ context.Context, flag string) (bool, error) {
1226+
return flag == outputSchemasFeatureFlag, nil
1227+
},
1228+
wantOutputSchema: true,
1229+
},
1230+
}
1231+
1232+
for _, tt := range tests {
1233+
t.Run(tt.name, func(t *testing.T) {
1234+
reg := mustBuild(t, NewBuilder().SetTools(tools).WithToolsets([]string{"all"}).WithFeatureChecker(tt.featureChecker))
1235+
1236+
registeredTools := listRegisteredTools(context.Background(), t, reg)
1237+
require.Len(t, registeredTools, 1)
1238+
1239+
if tt.wantOutputSchema {
1240+
require.NotNil(t, registeredTools[0].OutputSchema)
1241+
requireJSONEqual(t, outputSchema, registeredTools[0].OutputSchema)
1242+
} else {
1243+
require.Nil(t, registeredTools[0].OutputSchema)
1244+
}
1245+
1246+
require.Nil(t, reg.AllTools()[0].Tool.OutputSchema)
1247+
})
1248+
}
1249+
}
1250+
1251+
func listRegisteredTools(ctx context.Context, t *testing.T, reg *Inventory) []*mcp.Tool {
1252+
t.Helper()
1253+
1254+
srv := mcp.NewServer(&mcp.Implementation{Name: "test-server"}, nil)
1255+
reg.RegisterTools(ctx, srv, nil)
1256+
1257+
st, ct := mcp.NewInMemoryTransports()
1258+
client := mcp.NewClient(&mcp.Implementation{Name: "test-client"}, nil)
1259+
1260+
type connectResult struct {
1261+
session *mcp.ClientSession
1262+
err error
1263+
}
1264+
connectCh := make(chan connectResult, 1)
1265+
go func() {
1266+
clientSession, err := client.Connect(ctx, ct, nil)
1267+
connectCh <- connectResult{session: clientSession, err: err}
1268+
}()
1269+
1270+
serverSession, err := srv.Connect(ctx, st, nil)
1271+
require.NoError(t, err)
1272+
t.Cleanup(func() { require.NoError(t, serverSession.Close()) })
1273+
1274+
clientResult := <-connectCh
1275+
require.NoError(t, clientResult.err)
1276+
require.NotNil(t, clientResult.session)
1277+
t.Cleanup(func() { require.NoError(t, clientResult.session.Close()) })
1278+
1279+
result, err := clientResult.session.ListTools(ctx, nil)
1280+
require.NoError(t, err)
1281+
return result.Tools
1282+
}
1283+
1284+
func requireJSONEqual(t *testing.T, expected, actual any) {
1285+
t.Helper()
1286+
1287+
expectedJSON, err := json.Marshal(expected)
1288+
require.NoError(t, err)
1289+
actualJSON, err := json.Marshal(actual)
1290+
require.NoError(t, err)
1291+
require.JSONEq(t, string(expectedJSON), string(actualJSON))
1292+
}
1293+
12031294
func TestFeatureFlagResources(t *testing.T) {
12041295
resources := []ServerResourceTemplate{
12051296
mockResource("always_available", "toolset1", "uri1"),

0 commit comments

Comments
 (0)