Skip to content

Commit caa7792

Browse files
nlarewnlarew
committed
Update tests for custom fetch + untrusted content formatting
1 parent 4d887c9 commit caa7792

File tree

3 files changed

+100
-33
lines changed

3 files changed

+100
-33
lines changed

‎tests/integration/tools/assistant/assistantHelpers.ts‎

Lines changed: 11 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -1,6 +1,6 @@
11
import { setupIntegrationTest, IntegrationTest, defaultTestConfig, defaultDriverOptions } from "../../helpers.js";
22
import { describe, SuiteCollector } from "vitest";
3-
import { vi, beforeAll, afterAll } from "vitest";
3+
import { vi, beforeAll, afterAll, beforeEach } from "vitest";
44

55
export type IntegrationTestFunction = (integration: IntegrationTest) => void;
66

@@ -9,7 +9,7 @@ export function describeWithAssistant(name: string, fn: IntegrationTestFunction)
99
const integration = setupIntegrationTest(
1010
() => ({
1111
...defaultTestConfig,
12-
assistantBaseUrl: "https://knowledge.test.mongodb.com/api/", // Use test URL
12+
assistantBaseUrl: "https://knowledge.test.mongodb.com/api/", // Not a real URL
1313
}),
1414
() => ({
1515
...defaultDriverOptions,
@@ -39,8 +39,13 @@ interface MockedAssistantAPI {
3939
export function makeMockAssistantAPI(): MockedAssistantAPI {
4040
const mockFetch = vi.fn();
4141

42-
beforeAll(() => {
43-
global.fetch = mockFetch;
42+
beforeAll(async () => {
43+
const { createFetch } = await import("@mongodb-js/devtools-proxy-support");
44+
vi.mocked(createFetch).mockReturnValue(mockFetch as never);
45+
});
46+
47+
beforeEach(() => {
48+
mockFetch.mockClear();
4449
});
4550

4651
afterAll(() => {
@@ -50,14 +55,14 @@ export function makeMockAssistantAPI(): MockedAssistantAPI {
5055
const mockListSources: MockedAssistantAPI["mockListSources"] = (sources) => {
5156
mockFetch.mockResolvedValueOnce({
5257
ok: true,
53-
json: () => Promise.resolve({ dataSources: sources }),
58+
json: async () => ({ dataSources: sources }),
5459
});
5560
};
5661

5762
const mockSearchResults: MockedAssistantAPI["mockSearchResults"] = (results) => {
5863
mockFetch.mockResolvedValueOnce({
5964
ok: true,
60-
json: () => Promise.resolve({ results }),
65+
json: async () => ({ results }),
6166
});
6267
};
6368

‎tests/integration/tools/assistant/listKnowledgeSources.test.ts‎

Lines changed: 31 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -1,7 +1,18 @@
1-
import { describe, expect, it } from "vitest";
1+
import { describe, expect, it, vi } from "vitest";
22
import { CallToolResult } from "@modelcontextprotocol/sdk/types.js";
3-
import { expectDefined, validateToolMetadata, getResponseElements } from "../../helpers.js";
3+
import {
4+
expectDefined,
5+
validateToolMetadata,
6+
getResponseElements,
7+
getDataFromUntrustedContent,
8+
} from "../../helpers.js";
49
import { describeWithAssistant, makeMockAssistantAPI } from "./assistantHelpers.js";
10+
import { parse as yamlParse } from "yaml";
11+
12+
// Mock the devtools-proxy-support module
13+
vi.mock("@mongodb-js/devtools-proxy-support", () => ({
14+
createFetch: vi.fn(),
15+
}));
516

617
describeWithAssistant("list-knowledge-sources", (integration) => {
718
const { mockListSources, mockAPIError, mockNetworkError } = makeMockAssistantAPI();
@@ -46,20 +57,30 @@ describeWithAssistant("list-knowledge-sources", (integration) => {
4657

4758
const elements = getResponseElements(response.content);
4859

60+
// First element is the description
61+
expect(elements[0]?.text).toBe("Found 2 data sources in the MongoDB Assistant knowledge base.");
62+
63+
// Second element contains the YAML data
64+
expect(elements[1]?.text).toContain("<untrusted-user-data-");
65+
const yamlData = getDataFromUntrustedContent(elements[1]?.text ?? "");
66+
const dataSources = yamlParse(yamlData);
67+
4968
// Check first data source
50-
expect(elements[0]?.text).toBe("mongodb-manual");
51-
expect(elements[0]?._meta).toEqual({
69+
expect(dataSources[0]).toMatchObject({
70+
id: "mongodb-manual",
5271
type: "documentation",
72+
currentVersion: "7.0",
5373
versions: [
5474
{ label: "7.0", isCurrent: true },
5575
{ label: "6.0", isCurrent: false },
5676
],
5777
});
5878

5979
// Check second data source
60-
expect(elements[1]?.text).toBe("node-driver");
61-
expect(elements[1]?._meta).toEqual({
80+
expect(dataSources[1]).toMatchObject({
81+
id: "node-driver",
6282
type: "driver",
83+
currentVersion: "6.0",
6384
versions: [
6485
{ label: "6.0", isCurrent: true },
6586
{ label: "5.0", isCurrent: false },
@@ -76,7 +97,10 @@ describeWithAssistant("list-knowledge-sources", (integration) => {
7697

7798
expect(response.isError).toBeFalsy();
7899
expect(response.content).toBeInstanceOf(Array);
79-
expect(response.content).toHaveLength(0);
100+
expect(response.content).toHaveLength(2);
101+
102+
const elements = getResponseElements(response.content);
103+
expect(elements[0]?.text).toBe("Found 0 data sources in the MongoDB Assistant knowledge base.");
80104
});
81105
});
82106

‎tests/integration/tools/assistant/searchKnowledge.test.ts‎

Lines changed: 58 additions & 20 deletions
Original file line numberDiff line numberDiff line change
@@ -1,12 +1,19 @@
1-
import { describe, expect, it } from "vitest";
1+
import { describe, expect, it, vi } from "vitest";
22
import { CallToolResult } from "@modelcontextprotocol/sdk/types.js";
33
import {
44
expectDefined,
55
validateToolMetadata,
66
validateThrowsForInvalidArguments,
77
getResponseElements,
8+
getDataFromUntrustedContent,
89
} from "../../helpers.js";
910
import { describeWithAssistant, makeMockAssistantAPI } from "./assistantHelpers.js";
11+
import { parse as yamlParse } from "yaml";
12+
13+
// Mock the devtools-proxy-support module
14+
vi.mock("@mongodb-js/devtools-proxy-support", () => ({
15+
createFetch: vi.fn(),
16+
}));
1017

1118
describeWithAssistant("search-knowledge", (integration) => {
1219
const { mockSearchResults, mockAPIError, mockNetworkError } = makeMockAssistantAPI();
@@ -84,20 +91,34 @@ describeWithAssistant("search-knowledge", (integration) => {
8491

8592
const elements = getResponseElements(response.content);
8693

94+
// First element is the description
95+
expect(elements[0]?.text).toBe("Found 2 results in the MongoDB Assistant knowledge base.");
96+
97+
// Second element contains the YAML data
98+
expect(elements[1]?.text).toContain("<untrusted-user-data-");
99+
const yamlData = getDataFromUntrustedContent(elements[1]?.text ?? "");
100+
const results = yamlParse(yamlData);
101+
87102
// Check first result
88-
expect(elements[0]?.text).toBe(
89-
"The aggregation pipeline is a framework for data aggregation modeled on the concept of data processing pipelines."
90-
);
91-
expect(elements[0]?._meta).toEqual({
92-
tags: ["aggregation", "pipeline"],
93-
source: "mongodb-manual",
103+
expect(results[0]).toMatchObject({
104+
url: "https://docs.mongodb.com/manual/aggregation/",
105+
title: "Aggregation Pipeline",
106+
text: "The aggregation pipeline is a framework for data aggregation modeled on the concept of data processing pipelines.",
107+
metadata: {
108+
tags: ["aggregation", "pipeline"],
109+
source: "mongodb-manual",
110+
},
94111
});
95112

96113
// Check second result
97-
expect(elements[1]?.text).toBe("Aggregation pipeline operations have an array of operators available.");
98-
expect(elements[1]?._meta).toEqual({
99-
tags: ["aggregation", "operators"],
100-
source: "mongodb-manual",
114+
expect(results[1]).toMatchObject({
115+
url: "https://docs.mongodb.com/manual/reference/operator/aggregation/",
116+
title: "Aggregation Pipeline Operators",
117+
text: "Aggregation pipeline operations have an array of operators available.",
118+
metadata: {
119+
tags: ["aggregation", "operators"],
120+
source: "mongodb-manual",
121+
},
101122
});
102123
});
103124

@@ -127,15 +148,22 @@ describeWithAssistant("search-knowledge", (integration) => {
127148

128149
expect(response.isError).toBeFalsy();
129150
expect(response.content).toBeInstanceOf(Array);
130-
expect(response.content).toHaveLength(1);
151+
expect(response.content).toHaveLength(2);
131152

132153
const elements = getResponseElements(response.content);
133-
expect(elements[0]?.text).toBe(
134-
"The official MongoDB driver for Node.js provides a high-level API on top of mongodb-core."
135-
);
136-
expect(elements[0]?._meta).toEqual({
137-
tags: ["driver", "nodejs"],
138-
source: "node-driver",
154+
expect(elements[0]?.text).toBe("Found 1 results in the MongoDB Assistant knowledge base.");
155+
156+
const yamlData = getDataFromUntrustedContent(elements[1]?.text ?? "");
157+
const results = yamlParse(yamlData);
158+
159+
expect(results[0]).toMatchObject({
160+
url: "https://mongodb.github.io/node-mongodb-native/",
161+
title: "Node.js Driver",
162+
text: "The official MongoDB driver for Node.js provides a high-level API on top of mongodb-core.",
163+
metadata: {
164+
tags: ["driver", "nodejs"],
165+
source: "node-driver",
166+
},
139167
});
140168
});
141169

@@ -148,7 +176,10 @@ describeWithAssistant("search-knowledge", (integration) => {
148176

149177
expect(response.isError).toBeFalsy();
150178
expect(response.content).toBeInstanceOf(Array);
151-
expect(response.content).toHaveLength(0);
179+
expect(response.content).toHaveLength(2);
180+
181+
const elements = getResponseElements(response.content);
182+
expect(elements[0]?.text).toBe("Found 0 results in the MongoDB Assistant knowledge base.");
152183
});
153184

154185
it("uses default limit when not specified", async () => {
@@ -168,7 +199,14 @@ describeWithAssistant("search-knowledge", (integration) => {
168199
.callTool({ name: "search-knowledge", arguments: { query: "test query" } })) as CallToolResult;
169200

170201
expect(response.isError).toBeFalsy();
171-
expect(response.content).toHaveLength(5);
202+
expect(response.content).toHaveLength(2);
203+
204+
const elements = getResponseElements(response.content);
205+
expect(elements[0]?.text).toBe("Found 5 results in the MongoDB Assistant knowledge base.");
206+
207+
const yamlData = getDataFromUntrustedContent(elements[1]?.text ?? "");
208+
const results = yamlParse(yamlData);
209+
expect(results).toHaveLength(5);
172210
});
173211
});
174212

0 commit comments

Comments
 (0)