Jules unit tests for auth mcp tools (#9031)

* Jules unit tests for auth mcp tools

* Formats

* Fix disable_user_spec file.

* Update get_user spec

* Update any to ServerToolContext

* Update set_claims spec

* Update set_sms_region_policy

* Fixing broken error behavior - thanks unit tests!

---------

Co-authored-by: Alexander Nohe <nohe@google.com>

Author: joehan

src/gcp/auth.ts

@@ -12,7 +12,7 @@ interface MfaEnrollment {
   unobfuscatedPhoneInfo?: string;
 }
 
-interface UserInfo {
+export interface UserInfo {
   uid?: string;
   localId?: string;
   email: string;

src/mcp/tools/auth/disable_user.spec.ts

@@ -0,0 +1,63 @@
+import { expect } from "chai";
+import * as sinon from "sinon";
+import { disable_user } from "./disable_user";
+import * as auth from "../../../gcp/auth";
+import { toContent } from "../../util";
+import { ServerToolContext } from "../../tool";
+
+describe("disable_user tool", () => {
+  const projectId = "test-project";
+  const uid = "test-uid";
+
+  let disableUserStub: sinon.SinonStub;
+
+  beforeEach(() => {
+    disableUserStub = sinon.stub(auth, "disableUser");
+  });
+
+  afterEach(() => {
+    sinon.restore();
+  });
+
+  it("should disable a user successfully", async () => {
+    disableUserStub.resolves(true);
+
+    const result = await disable_user.fn({ uid, disabled: true }, {
+      projectId,
+    } as ServerToolContext);
+
+    expect(disableUserStub).to.be.calledWith(projectId, uid, true);
+    expect(result).to.deep.equal(toContent(`User ${uid} has been disabled`));
+  });
+
+  it("should enable a user successfully", async () => {
+    disableUserStub.resolves(true);
+
+    const result = await disable_user.fn({ uid, disabled: false }, {
+      projectId,
+    } as ServerToolContext);
+
+    expect(disableUserStub).to.be.calledWith(projectId, uid, false);
+    expect(result).to.deep.equal(toContent(`User ${uid} has been enabled`));
+  });
+
+  it("should handle failure to disable a user", async () => {
+    disableUserStub.resolves(false);
+
+    const result = await disable_user.fn({ uid, disabled: true }, {
+      projectId,
+    } as ServerToolContext);
+
+    expect(result).to.deep.equal(toContent(`Failed to disable user ${uid}`));
+  });
+
+  it("should handle failure to enable a user", async () => {
+    disableUserStub.resolves(false);
+
+    const result = await disable_user.fn({ uid, disabled: false }, {
+      projectId,
+    } as ServerToolContext);
+
+    expect(result).to.deep.equal(toContent(`Failed to enable user ${uid}`));
+  });
+});

src/mcp/tools/auth/disable_user.ts

@@ -24,7 +24,7 @@ export const disable_user = tool(
   async ({ uid, disabled }, { projectId }) => {
     const res = await disableUser(projectId, uid, disabled);
     if (res) {
-      return toContent(`User ${uid} as been ${disabled ? "disabled" : "enabled"}`);
+      return toContent(`User ${uid} has been ${disabled ? "disabled" : "enabled"}`);
     }
     return toContent(`Failed to ${disabled ? "disable" : "enable"} user ${uid}`);
   },

src/mcp/tools/auth/get_user.spec.ts

@@ -0,0 +1,60 @@
+import { expect } from "chai";
+import * as sinon from "sinon";
+import { get_user } from "./get_user";
+import * as auth from "../../../gcp/auth";
+import * as util from "../../util";
+import { ServerToolContext } from "../../tool";
+
+describe("get_user tool", () => {
+  const projectId = "test-project";
+  const email = "test@example.com";
+  const phoneNumber = "+11234567890";
+  const uid = "test-uid";
+  const user = { uid, email, phoneNumber };
+
+  let findUserStub: sinon.SinonStub;
+  let mcpErrorStub: sinon.SinonStub;
+
+  beforeEach(() => {
+    findUserStub = sinon.stub(auth, "findUser");
+    mcpErrorStub = sinon.stub(util, "mcpError");
+  });
+
+  afterEach(() => {
+    sinon.restore();
+  });
+
+  it("should return an error if no identifier is provided", async () => {
+    await get_user.fn({}, { projectId } as ServerToolContext);
+    expect(mcpErrorStub).to.be.calledWith("No user identifier supplied in auth_get_user tool");
+  });
+
+  it("should get a user by email", async () => {
+    findUserStub.resolves(user);
+    const result = await get_user.fn({ email }, { projectId } as ServerToolContext);
+    expect(findUserStub).to.be.calledWith(projectId, email, undefined, undefined);
+    expect(result).to.deep.equal(util.toContent(user));
+  });
+
+  it("should get a user by phone number", async () => {
+    findUserStub.resolves(user);
+    const result = await get_user.fn({ phone_number: phoneNumber }, {
+      projectId,
+    } as ServerToolContext);
+    expect(findUserStub).to.be.calledWith(projectId, undefined, phoneNumber, undefined);
+    expect(result).to.deep.equal(util.toContent(user));
+  });
+
+  it("should get a user by UID", async () => {
+    findUserStub.resolves(user);
+    const result = await get_user.fn({ uid }, { projectId } as ServerToolContext);
+    expect(findUserStub).to.be.calledWith(projectId, undefined, undefined, uid);
+    expect(result).to.deep.equal(util.toContent(user));
+  });
+
+  it("returns an error when no user exists", async () => {
+    findUserStub.rejects(new Error("No users found"));
+    await get_user.fn({ uid: "nonexistant@email.com" }, { projectId } as ServerToolContext);
+    expect(mcpErrorStub).to.be.calledWith("Unable to find user");
+  });
+});

src/mcp/tools/auth/get_user.ts

@@ -1,7 +1,7 @@
 import { z } from "zod";
 import { tool } from "../../tool";
 import { mcpError, toContent } from "../../util";
-import { findUser } from "../../../gcp/auth";
+import { findUser, UserInfo } from "../../../gcp/auth";
 
 export const get_user = tool(
   {
@@ -36,8 +36,14 @@ export const get_user = tool(
   },
   async ({ email, phone_number, uid }, { projectId }) => {
     if (email === undefined && phone_number === undefined && uid === undefined) {
-      return mcpError(`No user identifier supplied in auth_get_user tool`);
+      return mcpError("No user identifier supplied in auth_get_user tool");
     }
-    return toContent(await findUser(projectId, email, phone_number, uid));
+    let user: UserInfo;
+    try {
+      user = await findUser(projectId, email, phone_number, uid);
+    } catch (err: any) {
+      return mcpError("Unable to find user");
+    }
+    return toContent(user);
   },
 );

src/mcp/tools/auth/list_users.spec.ts

@@ -0,0 +1,55 @@
+import { expect } from "chai";
+import * as sinon from "sinon";
+import { list_users } from "./list_users";
+import * as auth from "../../../gcp/auth";
+import { toContent } from "../../util";
+import { ServerToolContext } from "../../tool";
+
+describe("list_users tool", () => {
+  const projectId = "test-project";
+  const users = [
+    { uid: "uid1", email: "user1@example.com", passwordHash: "hash", salt: "salt" },
+    { uid: "uid2", email: "user2@example.com", passwordHash: "hash", salt: "salt" },
+  ];
+  const prunedUsers = [
+    { uid: "uid1", email: "user1@example.com" },
+    { uid: "uid2", email: "user2@example.com" },
+  ];
+
+  let listUsersStub: sinon.SinonStub;
+
+  beforeEach(() => {
+    listUsersStub = sinon.stub(auth, "listUsers");
+  });
+
+  afterEach(() => {
+    sinon.restore();
+  });
+
+  it("should list users with the default limit", async () => {
+    listUsersStub.resolves(users);
+
+    const result = await list_users.fn({}, { projectId } as ServerToolContext);
+
+    expect(listUsersStub).to.be.calledWith(projectId, 100);
+    expect(result).to.deep.equal(toContent(prunedUsers));
+  });
+
+  it("should list users with a specified limit", async () => {
+    listUsersStub.resolves(users);
+
+    const result = await list_users.fn({ limit: 10 }, { projectId } as ServerToolContext);
+
+    expect(listUsersStub).to.be.calledWith(projectId, 10);
+    expect(result).to.deep.equal(toContent(prunedUsers));
+  });
+
+  it("should handle an empty list of users", async () => {
+    listUsersStub.resolves([]);
+
+    const result = await list_users.fn({}, { projectId } as ServerToolContext);
+
+    expect(listUsersStub).to.be.calledWith(projectId, 100);
+    expect(result).to.deep.equal(toContent([]));
+  });
+});

src/mcp/tools/auth/set_claims.spec.ts

@@ -0,0 +1,72 @@
+import { expect } from "chai";
+import * as sinon from "sinon";
+import { set_claim } from "./set_claims";
+import * as auth from "../../../gcp/auth";
+import * as util from "../../util";
+import { ServerToolContext } from "../../tool";
+
+describe("set_claim tool", () => {
+  const projectId = "test-project";
+  const uid = "test-uid";
+  const claim = "admin";
+
+  let setCustomClaimStub: sinon.SinonStub;
+  let mcpErrorStub: sinon.SinonStub;
+
+  beforeEach(() => {
+    setCustomClaimStub = sinon.stub(auth, "setCustomClaim");
+    mcpErrorStub = sinon.stub(util, "mcpError");
+  });
+
+  afterEach(() => {
+    sinon.restore();
+  });
+
+  it("should set a simple claim", async () => {
+    const value = true;
+    setCustomClaimStub.resolves({ success: true });
+
+    const result = await set_claim.fn({ uid, claim, value }, { projectId } as ServerToolContext);
+
+    expect(setCustomClaimStub).to.be.calledWith(
+      projectId,
+      uid,
+      { [claim]: value },
+      { merge: true },
+    );
+    expect(result).to.deep.equal(util.toContent({ success: true }));
+  });
+
+  it("should set a JSON claim", async () => {
+    const json_value = '{"role": "editor"}';
+    const parsedValue = { role: "editor" };
+    setCustomClaimStub.resolves({ success: true });
+
+    const result = await set_claim.fn({ uid, claim, json_value }, {
+      projectId,
+    } as ServerToolContext);
+
+    expect(setCustomClaimStub).to.be.calledWith(
+      projectId,
+      uid,
+      { [claim]: parsedValue },
+      { merge: true },
+    );
+    expect(result).to.deep.equal(util.toContent({ success: true }));
+  });
+
+  it("should return an error for invalid JSON", async () => {
+    const json_value = "invalid-json";
+    await set_claim.fn({ uid, claim, json_value }, { projectId } as ServerToolContext);
+    expect(mcpErrorStub).to.be.calledWith(
+      `Provided \`json_value\` was not valid JSON: ${json_value}`,
+    );
+  });
+
+  it("should return an error if both value and json_value are provided", async () => {
+    const value = "simple";
+    const json_value = '{"complex": true}';
+    await set_claim.fn({ uid, claim, value, json_value }, { projectId } as ServerToolContext);
+    expect(mcpErrorStub).to.be.calledWith("Must supply only `value` or `json_value`, not both.");
+  });
+});

src/mcp/tools/auth/set_sms_region_policy.spec.ts

@@ -0,0 +1,48 @@
+import { expect } from "chai";
+import * as sinon from "sinon";
+import { set_sms_region_policy } from "./set_sms_region_policy";
+import * as auth from "../../../gcp/auth";
+import { toContent } from "../../util";
+import { ServerToolContext } from "../../tool";
+
+describe("set_sms_region_policy tool", () => {
+  const projectId = "test-project";
+  const country_codes = ["us", "ca"];
+  const upperCaseCountryCodes = ["US", "CA"];
+
+  let setAllowSmsRegionPolicyStub: sinon.SinonStub;
+  let setDenySmsRegionPolicyStub: sinon.SinonStub;
+
+  beforeEach(() => {
+    setAllowSmsRegionPolicyStub = sinon.stub(auth, "setAllowSmsRegionPolicy");
+    setDenySmsRegionPolicyStub = sinon.stub(auth, "setDenySmsRegionPolicy");
+  });
+
+  afterEach(() => {
+    sinon.restore();
+  });
+
+  it("should set an ALLOW policy", async () => {
+    setAllowSmsRegionPolicyStub.resolves(true);
+
+    const result = await set_sms_region_policy.fn({ policy_type: "ALLOW", country_codes }, {
+      projectId,
+    } as ServerToolContext);
+
+    expect(setAllowSmsRegionPolicyStub).to.be.calledWith(projectId, upperCaseCountryCodes);
+    expect(setDenySmsRegionPolicyStub).to.not.be.called;
+    expect(result).to.deep.equal(toContent(true));
+  });
+
+  it("should set a DENY policy", async () => {
+    setDenySmsRegionPolicyStub.resolves(true);
+
+    const result = await set_sms_region_policy.fn({ policy_type: "DENY", country_codes }, {
+      projectId,
+    } as ServerToolContext);
+
+    expect(setDenySmsRegionPolicyStub).to.be.calledWith(projectId, upperCaseCountryCodes);
+    expect(setAllowSmsRegionPolicyStub).to.not.be.called;
+    expect(result).to.deep.equal(toContent(true));
+  });
+});