|
| 1 | +package middleware |
| 2 | + |
| 3 | +// Tests targeting previously uncovered branches in jqschema.go: |
| 4 | +// - inferSchema default case (custom types via reflect) |
| 5 | +// - savePayload error paths (unwritable directories) |
| 6 | +// - WrapToolHandler with custom struct data (triggers default branch in type switch) |
| 7 | +// - WrapToolHandler when savePayload fails (continues and returns metadata on save failure) |
| 8 | + |
| 9 | +import ( |
| 10 | + "context" |
| 11 | + "encoding/json" |
| 12 | + "os" |
| 13 | + "path/filepath" |
| 14 | + "runtime" |
| 15 | + "testing" |
| 16 | + |
| 17 | + sdk "github.com/modelcontextprotocol/go-sdk/mcp" |
| 18 | + "github.com/stretchr/testify/assert" |
| 19 | + "github.com/stretchr/testify/require" |
| 20 | +) |
| 21 | + |
| 22 | +// --------------------------------------------------------------------------- |
| 23 | +// inferSchema default case (lines 104-115 in jqschema.go) |
| 24 | +// --------------------------------------------------------------------------- |
| 25 | + |
| 26 | +// customIntType has kind reflect.Int32, which is NOT listed in the explicit |
| 27 | +// type-switch cases (those only list the untyped int/int32 etc. kinds), so it |
| 28 | +// falls through to the default branch and is classified via reflect. |
| 29 | +type customIntType int32 |
| 30 | + |
| 31 | +// customUintType similarly exercises the uint reflect path. |
| 32 | +type customUintType uint64 |
| 33 | + |
| 34 | +// customFloatType exercises the float reflect path. |
| 35 | +type customFloatType float32 |
| 36 | + |
| 37 | +// myStruct is an arbitrary struct that falls through to the inner "string" default. |
| 38 | +type myStruct struct { |
| 39 | + X int |
| 40 | + Y string |
| 41 | +} |
| 42 | + |
| 43 | +func TestInferSchema_DefaultCaseFallback(t *testing.T) { |
| 44 | + tests := []struct { |
| 45 | + name string |
| 46 | + input interface{} |
| 47 | + expected interface{} |
| 48 | + }{ |
| 49 | + { |
| 50 | + name: "custom named int type returns number", |
| 51 | + input: customIntType(42), |
| 52 | + expected: "number", |
| 53 | + }, |
| 54 | + { |
| 55 | + name: "custom named uint type returns number", |
| 56 | + input: customUintType(100), |
| 57 | + expected: "number", |
| 58 | + }, |
| 59 | + { |
| 60 | + name: "custom named float type returns number", |
| 61 | + input: customFloatType(3.14), |
| 62 | + expected: "number", |
| 63 | + }, |
| 64 | + { |
| 65 | + name: "struct type returns string", |
| 66 | + input: myStruct{X: 1, Y: "hello"}, |
| 67 | + expected: "string", |
| 68 | + }, |
| 69 | + { |
| 70 | + name: "chan type returns string", |
| 71 | + input: make(chan int), |
| 72 | + expected: "string", |
| 73 | + }, |
| 74 | + { |
| 75 | + // []int is NOT []interface{} so it falls through to default. |
| 76 | + name: "typed int slice returns string", |
| 77 | + input: []int{1, 2, 3}, |
| 78 | + expected: "string", |
| 79 | + }, |
| 80 | + { |
| 81 | + // map[string]string is NOT map[string]interface{} so it falls through. |
| 82 | + name: "typed string map returns string", |
| 83 | + input: map[string]string{"key": "value"}, |
| 84 | + expected: "string", |
| 85 | + }, |
| 86 | + } |
| 87 | + |
| 88 | + for _, tt := range tests { |
| 89 | + t.Run(tt.name, func(t *testing.T) { |
| 90 | + got := inferSchema(tt.input) |
| 91 | + assert.Equal(t, tt.expected, got) |
| 92 | + }) |
| 93 | + } |
| 94 | +} |
| 95 | + |
| 96 | +// TestInferSchema_DefaultCaseNestedInMap exercises the default branch when a |
| 97 | +// custom type appears as a value inside a map[string]interface{}. |
| 98 | +func TestInferSchema_DefaultCaseNestedInMap(t *testing.T) { |
| 99 | + input := map[string]interface{}{ |
| 100 | + "normal": float64(1), |
| 101 | + "custom": customIntType(99), |
| 102 | + "strType": []int{1, 2, 3}, |
| 103 | + } |
| 104 | + got := inferSchema(input) |
| 105 | + schema, ok := got.(map[string]interface{}) |
| 106 | + require.True(t, ok) |
| 107 | + assert.Equal(t, "number", schema["normal"]) |
| 108 | + assert.Equal(t, "number", schema["custom"]) // customIntType → reflect.Int32 → "number" |
| 109 | + assert.Equal(t, "string", schema["strType"]) // []int → reflect.Slice → "string" |
| 110 | +} |
| 111 | + |
| 112 | +// --------------------------------------------------------------------------- |
| 113 | +// savePayload error paths (lines 231-256 in jqschema.go) |
| 114 | +// --------------------------------------------------------------------------- |
| 115 | + |
| 116 | +// TestSavePayload_MkdirAllFailure covers the os.MkdirAll error path by using a |
| 117 | +// base directory whose parent is read-only (preventing subdirectory creation). |
| 118 | +func TestSavePayload_MkdirAllFailure(t *testing.T) { |
| 119 | + if runtime.GOOS == "windows" { |
| 120 | + t.Skip("permission model differs on Windows") |
| 121 | + } |
| 122 | + |
| 123 | + // Create a temporary directory and make it read-only so that MkdirAll |
| 124 | + // cannot create any subdirectories inside it. |
| 125 | + readOnlyDir := t.TempDir() |
| 126 | + require.NoError(t, os.Chmod(readOnlyDir, 0444)) |
| 127 | + t.Cleanup(func() { |
| 128 | + // Restore permissions so t.TempDir cleanup can remove the directory. |
| 129 | + _ = os.Chmod(readOnlyDir, 0755) |
| 130 | + }) |
| 131 | + |
| 132 | + // Skip the test when running as root (root can write to read-only dirs). |
| 133 | + if os.Getuid() == 0 { |
| 134 | + t.Skip("skipping read-only directory test when running as root") |
| 135 | + } |
| 136 | + |
| 137 | + _, err := savePayload(readOnlyDir, "", "session1", "query1", []byte(`{"key":"value"}`)) |
| 138 | + require.Error(t, err, "savePayload should return an error when MkdirAll fails") |
| 139 | + assert.Contains(t, err.Error(), "failed to create payload directory") |
| 140 | +} |
| 141 | + |
| 142 | +// TestSavePayload_WriteFileFailure covers the os.WriteFile error path by using a |
| 143 | +// directory path that resolves to an existing file (so the path cannot be a dir). |
| 144 | +func TestSavePayload_WriteFileFailure(t *testing.T) { |
| 145 | + if runtime.GOOS == "windows" { |
| 146 | + t.Skip("permission model differs on Windows") |
| 147 | + } |
| 148 | + if os.Getuid() == 0 { |
| 149 | + t.Skip("skipping file-permission test when running as root") |
| 150 | + } |
| 151 | + |
| 152 | + // Create the target directory structure that savePayload would create, |
| 153 | + // then place a read-only file at the exact payload path. |
| 154 | + baseDir := t.TempDir() |
| 155 | + sessionID := "test-session" |
| 156 | + queryID := "test-query" |
| 157 | + payloadDir := filepath.Join(baseDir, sessionID, queryID) |
| 158 | + require.NoError(t, os.MkdirAll(payloadDir, 0755)) |
| 159 | + |
| 160 | + // Write a read-only file at the exact location savePayload would write. |
| 161 | + payloadFile := filepath.Join(payloadDir, "payload.json") |
| 162 | + require.NoError(t, os.WriteFile(payloadFile, []byte("existing"), 0444)) |
| 163 | + t.Cleanup(func() { _ = os.Chmod(payloadFile, 0644) }) |
| 164 | + |
| 165 | + _, err := savePayload(baseDir, "", sessionID, queryID, []byte(`{"key":"value"}`)) |
| 166 | + // On Linux, writing to a 0444 file returns "permission denied" |
| 167 | + require.Error(t, err, "savePayload should return an error when WriteFile fails") |
| 168 | + assert.Contains(t, err.Error(), "failed to write payload file") |
| 169 | +} |
| 170 | + |
| 171 | +// --------------------------------------------------------------------------- |
| 172 | +// WrapToolHandler with custom struct data (default branch in type switch) |
| 173 | +// --------------------------------------------------------------------------- |
| 174 | + |
| 175 | +// TestWrapToolHandler_CustomStructData exercises the default arm of the data |
| 176 | +// type switch inside WrapToolHandler (lines 388-391 in jqschema.go). A plain |
| 177 | +// Go struct is returned as the data value so that none of the explicit cases |
| 178 | +// (map[string]interface{}, []interface{}, string, float64, bool, nil, |
| 179 | +// json.Number) match. |
| 180 | +func TestWrapToolHandler_CustomStructData(t *testing.T) { |
| 181 | + baseDir := t.TempDir() |
| 182 | + |
| 183 | + type ToolOutput struct { |
| 184 | + Count int `json:"count"` |
| 185 | + Label string `json:"label"` |
| 186 | + } |
| 187 | + |
| 188 | + // Return a custom struct as data. |
| 189 | + customData := ToolOutput{Count: 42, Label: "hello"} |
| 190 | + mockHandler := func(ctx context.Context, req *sdk.CallToolRequest, args interface{}) (*sdk.CallToolResult, interface{}, error) { |
| 191 | + return &sdk.CallToolResult{IsError: false}, customData, nil |
| 192 | + } |
| 193 | + |
| 194 | + // Use threshold 0 so the large-payload path is always taken. |
| 195 | + wrapped := WrapToolHandler(mockHandler, "test_tool", baseDir, "", 0, testGetSessionID) |
| 196 | + result, _, err := wrapped(context.Background(), &sdk.CallToolRequest{}, nil) |
| 197 | + |
| 198 | + require.NoError(t, err) |
| 199 | + require.NotNil(t, result) |
| 200 | + assert.False(t, result.IsError) |
| 201 | + |
| 202 | + // Verify the response is a metadata object (large-payload path taken). |
| 203 | + require.NotEmpty(t, result.Content) |
| 204 | + textContent, ok := result.Content[0].(*sdk.TextContent) |
| 205 | + require.True(t, ok) |
| 206 | + |
| 207 | + var meta map[string]interface{} |
| 208 | + require.NoError(t, json.Unmarshal([]byte(textContent.Text), &meta)) |
| 209 | + // payloadSchema for {"count": 42, "label": "hello"} should be a map with "number" and "string". |
| 210 | + schema, ok := meta["payloadSchema"].(map[string]interface{}) |
| 211 | + require.True(t, ok, "payloadSchema should be an object, got: %T %v", meta["payloadSchema"], meta["payloadSchema"]) |
| 212 | + assert.Equal(t, "number", schema["count"]) |
| 213 | + assert.Equal(t, "string", schema["label"]) |
| 214 | +} |
| 215 | + |
| 216 | +// --------------------------------------------------------------------------- |
| 217 | +// WrapToolHandler when savePayload fails |
| 218 | +// (lines 358-365: save error is logged but processing continues) |
| 219 | +// --------------------------------------------------------------------------- |
| 220 | + |
| 221 | +// TestWrapToolHandler_SavePayloadFailure verifies that when savePayload cannot |
| 222 | +// write the file (e.g. unwritable base directory), WrapToolHandler logs the |
| 223 | +// error but still returns a metadata response (schema + preview) rather than |
| 224 | +// propagating the failure to the caller. |
| 225 | +func TestWrapToolHandler_SavePayloadFailure(t *testing.T) { |
| 226 | + if runtime.GOOS == "windows" { |
| 227 | + t.Skip("permission model differs on Windows") |
| 228 | + } |
| 229 | + if os.Getuid() == 0 { |
| 230 | + t.Skip("skipping read-only directory test when running as root") |
| 231 | + } |
| 232 | + |
| 233 | + // Make baseDir read-only so that MkdirAll cannot create subdirectories. |
| 234 | + baseDir := t.TempDir() |
| 235 | + require.NoError(t, os.Chmod(baseDir, 0444)) |
| 236 | + t.Cleanup(func() { _ = os.Chmod(baseDir, 0755) }) |
| 237 | + |
| 238 | + largeData := map[string]interface{}{ |
| 239 | + "items": make([]interface{}, 100), |
| 240 | + } |
| 241 | + for i := range largeData["items"].([]interface{}) { |
| 242 | + largeData["items"].([]interface{})[i] = map[string]interface{}{"id": float64(i)} |
| 243 | + } |
| 244 | + |
| 245 | + mockHandler := func(ctx context.Context, req *sdk.CallToolRequest, args interface{}) (*sdk.CallToolResult, interface{}, error) { |
| 246 | + return &sdk.CallToolResult{IsError: false}, largeData, nil |
| 247 | + } |
| 248 | + |
| 249 | + // Threshold of 1 forces large-payload path. |
| 250 | + wrapped := WrapToolHandler(mockHandler, "test_tool", baseDir, "", 1, testGetSessionID) |
| 251 | + result, _, err := wrapped(context.Background(), &sdk.CallToolRequest{}, nil) |
| 252 | + |
| 253 | + // WrapToolHandler must not propagate the savePayload error. |
| 254 | + require.NoError(t, err, "WrapToolHandler must not return an error when savePayload fails") |
| 255 | + require.NotNil(t, result) |
| 256 | + |
| 257 | + // Assert the observable failure-specific behavior: the wrapper still returns |
| 258 | + // metadata (schema + preview), but no payload path/file is produced. |
| 259 | + require.Len(t, result.Content, 1) |
| 260 | + textContent, ok := result.Content[0].(*sdk.TextContent) |
| 261 | + require.True(t, ok, "expected TextContent in result") |
| 262 | + |
| 263 | + var meta map[string]interface{} |
| 264 | + require.NoError(t, json.Unmarshal([]byte(textContent.Text), &meta)) |
| 265 | + assert.Contains(t, meta, "payloadSchema", "expected metadata response to include payloadSchema when payload saving fails") |
| 266 | + assert.Contains(t, meta, "payloadPreview", "expected metadata response to include payloadPreview when payload saving fails") |
| 267 | + // payloadPath should be empty (no file was saved). |
| 268 | + assert.Equal(t, "", meta["payloadPath"], "payloadPath should be empty when savePayload fails") |
| 269 | + |
| 270 | + entries, readErr := os.ReadDir(baseDir) |
| 271 | + require.NoError(t, readErr) |
| 272 | + assert.Len(t, entries, 0, "savePayload failure path should not leave any saved payload files or directories behind") |
| 273 | +} |
0 commit comments