Skip to content

chore: add integration tests for create-index #84

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
merged 5 commits into from
Apr 22, 2025
Merged
Show file tree
Hide file tree
Changes from 4 commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
11 changes: 9 additions & 2 deletions src/tools/mongodb/create/createIndex.ts
Original file line number Diff line number Diff line change
Expand Up @@ -10,22 +10,29 @@ export class CreateIndexTool extends MongoDBToolBase {
protected argsShape = {
...DbOperationArgs,
keys: z.record(z.string(), z.custom<IndexDirection>()).describe("The index definition"),
name: z.string().optional().describe("The name of the index"),
};

protected operationType: OperationType = "create";

protected async execute({ database, collection, keys }: ToolArgs<typeof this.argsShape>): Promise<CallToolResult> {
protected async execute({
database,
collection,
keys,
name,
}: ToolArgs<typeof this.argsShape>): Promise<CallToolResult> {
const provider = await this.ensureConnected();
const indexes = await provider.createIndexes(database, collection, [
{
key: keys,
name,
},
]);

return {
content: [
{
text: `Created the index \`${indexes[0]}\``,
text: `Created the index "${indexes[0]}" on collection "${collection}" in database "${database}"`,
type: "text",
},
],
Expand Down
33 changes: 14 additions & 19 deletions src/tools/mongodb/metadata/connect.ts
Original file line number Diff line number Diff line change
Expand Up @@ -2,7 +2,6 @@ import { z } from "zod";
import { CallToolResult } from "@modelcontextprotocol/sdk/types.js";
import { MongoDBToolBase } from "../mongodbTool.js";
import { ToolArgs, OperationType } from "../../tool.js";
import { ErrorCodes, MongoDBError } from "../../../errors.js";
import config from "../../../config.js";
import { MongoError as DriverError } from "mongodb";

Expand Down Expand Up @@ -37,25 +36,21 @@ export class ConnectTool extends MongoDBToolBase {

let connectionString: string;

if (typeof connectionStringOrClusterName === "string") {
if (
connectionStringOrClusterName.startsWith("mongodb://") ||
connectionStringOrClusterName.startsWith("mongodb+srv://")
) {
connectionString = connectionStringOrClusterName;
} else {
// TODO:
return {
content: [
{
type: "text",
text: `Connecting via cluster name not supported yet. Please provide a connection string.`,
},
],
};
}
if (
connectionStringOrClusterName.startsWith("mongodb://") ||
connectionStringOrClusterName.startsWith("mongodb+srv://")
) {
connectionString = connectionStringOrClusterName;
} else {
throw new MongoDBError(ErrorCodes.InvalidParams, "Invalid connection options");
// TODO:
Copy link
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

?

Copy link
Collaborator Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

This is unchanged code, but it's the hook where we'd want to call any atlas API that would allow us to generate a connection string based on a cluster name (if possible).

Copy link
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

should we add some text to that todo then?

Copy link
Collaborator Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

The text we return from the next statement is pretty much what we would write in the TODO comment, but happy to update it and add a link to the issue.

return {
content: [
{
type: "text",
text: `Connecting via cluster name not supported yet. Please provide a connection string.`,
},
],
};
}

try {
Expand Down
3 changes: 2 additions & 1 deletion src/tools/mongodb/mongodbTool.ts
Original file line number Diff line number Diff line change
Expand Up @@ -19,9 +19,10 @@ export abstract class MongoDBToolBase extends ToolBase {
protected category: ToolCategory = "mongodb";

protected async ensureConnected(): Promise<NodeDriverServiceProvider> {
const provider = this.session.serviceProvider;
let provider = this.session.serviceProvider;
Copy link
Collaborator

@gagik gagik Apr 22, 2025

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

maybe we should directly use this.session.serviceProvider at all times to avoid confusion

if (!provider && config.connectionString) {
await this.connectToMongoDB(config.connectionString);
provider = this.session.serviceProvider;
Copy link
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

It seems you're assigning the same value, would you double check if this is needed?

Copy link
Collaborator Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

connectToMongoDB - if successful - will set the session provider.

Copy link
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

is there a scenario where NodeDriverServiceProvider.connect returns undefined? seems like it'd throw instead right (and not be caught by connectToMongoDB)?

Copy link
Collaborator Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Based on the types, if the promise resolves, we should get a valid NodeServiceProvider and not undefined. If the connect attempt fails, this will throw an error which will eventually be bubbled up to the error handler of whatever tool the user was calling.

Copy link
Collaborator

@gagik gagik Apr 22, 2025

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

do I understand correctly then that the second if (!this.session.serviceProvider) would never occur? We could make the connectToMongoDB return the provider and set it to the local provider variable (I guess going back to that..) to make that type logic work

Copy link
Collaborator Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I changed it to always use this.session.serviceProvider and removed the local provider var.

}

if (!provider) {
Expand Down
10 changes: 9 additions & 1 deletion tests/integration/helpers.ts
Original file line number Diff line number Diff line change
Expand Up @@ -6,8 +6,9 @@ import path from "path";
import fs from "fs/promises";
import { Session } from "../../src/session.js";
import { McpServer } from "@modelcontextprotocol/sdk/server/mcp.js";
import { MongoClient } from "mongodb";
import { MongoClient, ObjectId } from "mongodb";
import { toIncludeAllMembers } from "jest-extended";
import config from "../../src/config.js";

interface ParameterInfo {
name: string;
Expand All @@ -23,13 +24,16 @@ export function setupIntegrationTest(): {
mongoClient: () => MongoClient;
connectionString: () => string;
connectMcpClient: () => Promise<void>;
randomDbName: () => string;
} {
let mongoCluster: runner.MongoCluster | undefined;
let mongoClient: MongoClient | undefined;

let mcpClient: Client | undefined;
let mcpServer: Server | undefined;

let randomDbName: string;

beforeEach(async () => {
const clientTransport = new InMemoryTransport();
const serverTransport = new InMemoryTransport();
Expand Down Expand Up @@ -59,6 +63,7 @@ export function setupIntegrationTest(): {
});
await mcpServer.connect(serverTransport);
await mcpClient.connect(clientTransport);
randomDbName = new ObjectId().toString();
});

afterEach(async () => {
Expand All @@ -70,6 +75,8 @@ export function setupIntegrationTest(): {

await mongoClient?.close();
mongoClient = undefined;

config.connectionString = undefined;
});

beforeAll(async function () {
Expand Down Expand Up @@ -144,6 +151,7 @@ export function setupIntegrationTest(): {
arguments: { connectionStringOrClusterName: getConnectionString() },
});
},
randomDbName: () => randomDbName,
};
}

Expand Down
64 changes: 43 additions & 21 deletions tests/integration/tools/mongodb/create/createCollection.test.ts
Original file line number Diff line number Diff line change
Expand Up @@ -7,6 +7,7 @@ import {
import { toIncludeSameMembers } from "jest-extended";
import { McpError } from "@modelcontextprotocol/sdk/types.js";
import { ObjectId } from "bson";
import config from "../../../../../src/config.js";

describe("createCollection tool", () => {
const integration = setupIntegrationTest();
Expand Down Expand Up @@ -48,69 +49,90 @@ describe("createCollection tool", () => {
describe("with non-existent database", () => {
it("creates a new collection", async () => {
const mongoClient = integration.mongoClient();
let collections = await mongoClient.db("foo").listCollections().toArray();
let collections = await mongoClient.db(integration.randomDbName()).listCollections().toArray();
expect(collections).toHaveLength(0);

await integration.connectMcpClient();
const response = await integration.mcpClient().callTool({
name: "create-collection",
arguments: { database: "foo", collection: "bar" },
arguments: { database: integration.randomDbName(), collection: "bar" },
});
const content = getResponseContent(response.content);
expect(content).toEqual('Collection "bar" created in database "foo".');
expect(content).toEqual(`Collection "bar" created in database "${integration.randomDbName()}".`);

collections = await mongoClient.db("foo").listCollections().toArray();
collections = await mongoClient.db(integration.randomDbName()).listCollections().toArray();
expect(collections).toHaveLength(1);
expect(collections[0].name).toEqual("bar");
});
});

describe("with existing database", () => {
let dbName: string;
beforeEach(() => {
dbName = new ObjectId().toString();
});

it("creates new collection", async () => {
const mongoClient = integration.mongoClient();
await mongoClient.db(dbName).createCollection("collection1");
let collections = await mongoClient.db(dbName).listCollections().toArray();
await mongoClient.db(integration.randomDbName()).createCollection("collection1");
let collections = await mongoClient.db(integration.randomDbName()).listCollections().toArray();
expect(collections).toHaveLength(1);

await integration.connectMcpClient();
const response = await integration.mcpClient().callTool({
name: "create-collection",
arguments: { database: dbName, collection: "collection2" },
arguments: { database: integration.randomDbName(), collection: "collection2" },
});
const content = getResponseContent(response.content);
expect(content).toEqual(`Collection "collection2" created in database "${dbName}".`);
collections = await mongoClient.db(dbName).listCollections().toArray();
expect(content).toEqual(`Collection "collection2" created in database "${integration.randomDbName()}".`);
collections = await mongoClient.db(integration.randomDbName()).listCollections().toArray();
expect(collections).toHaveLength(2);
expect(collections.map((c) => c.name)).toIncludeSameMembers(["collection1", "collection2"]);
});

it("does nothing if collection already exists", async () => {
const mongoClient = integration.mongoClient();
await mongoClient.db(dbName).collection("collection1").insertOne({});
let collections = await mongoClient.db(dbName).listCollections().toArray();
await mongoClient.db(integration.randomDbName()).collection("collection1").insertOne({});
let collections = await mongoClient.db(integration.randomDbName()).listCollections().toArray();
expect(collections).toHaveLength(1);
let documents = await mongoClient.db(dbName).collection("collection1").find({}).toArray();
let documents = await mongoClient
.db(integration.randomDbName())
.collection("collection1")
.find({})
.toArray();
expect(documents).toHaveLength(1);

await integration.connectMcpClient();
const response = await integration.mcpClient().callTool({
name: "create-collection",
arguments: { database: dbName, collection: "collection1" },
arguments: { database: integration.randomDbName(), collection: "collection1" },
});
const content = getResponseContent(response.content);
expect(content).toEqual(`Collection "collection1" created in database "${dbName}".`);
collections = await mongoClient.db(dbName).listCollections().toArray();
expect(content).toEqual(`Collection "collection1" created in database "${integration.randomDbName()}".`);
collections = await mongoClient.db(integration.randomDbName()).listCollections().toArray();
expect(collections).toHaveLength(1);
expect(collections[0].name).toEqual("collection1");

// Make sure we didn't drop the existing collection
documents = await mongoClient.db(dbName).collection("collection1").find({}).toArray();
documents = await mongoClient.db(integration.randomDbName()).collection("collection1").find({}).toArray();
expect(documents).toHaveLength(1);
});
});

describe("when not connected", () => {
it("connects automatically if connection string is configured", async () => {
config.connectionString = integration.connectionString();

const response = await integration.mcpClient().callTool({
name: "create-collection",
arguments: { database: integration.randomDbName(), collection: "new-collection" },
});
const content = getResponseContent(response.content);
expect(content).toEqual(`Collection "new-collection" created in database "${integration.randomDbName()}".`);
});

it("throw an error if connection string is not configured", async () => {
const response = await integration.mcpClient().callTool({
name: "create-collection",
arguments: { database: integration.randomDbName(), collection: "new-collection" },
});
const content = getResponseContent(response.content);
expect(content).toContain("You need to connect to a MongoDB instance before you can access its data.");
});
});
});
Loading
Loading