Skip to content

Commit db5059b

Browse files
taltultctaltultc
committed
feat(allocations): new allocations tools
1 parent 674b0c7 commit db5059b

File tree

7 files changed

+543
-3
lines changed

7 files changed

+543
-3
lines changed

‎doit-mcp-server/src/index.ts‎

Lines changed: 10 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -49,6 +49,12 @@ import {
4949
listInvoicesTool,
5050
getInvoiceTool,
5151
} from "../../src/tools/invoices.js";
52+
import {
53+
ListAllocationsArgumentsSchema,
54+
GetAllocationArgumentsSchema,
55+
listAllocationsTool,
56+
getAllocationTool,
57+
} from "../../src/tools/allocations.js";
5258
import {
5359
ChangeCustomerArgumentsSchema,
5460
changeCustomerTool,
@@ -250,6 +256,10 @@ export class DoitMCPAgent extends McpAgent {
250256
this.registerTool(listInvoicesTool, ListInvoicesArgumentsSchema);
251257
this.registerTool(getInvoiceTool, GetInvoiceArgumentsSchema);
252258

259+
// Allocations tools
260+
this.registerTool(listAllocationsTool, ListAllocationsArgumentsSchema);
261+
this.registerTool(getAllocationTool, GetAllocationArgumentsSchema);
262+
253263
// Change Customer tool (requires special handling)
254264
if (this.props.isDoitUser === "true") {
255265
(this.server.tool as any)(

‎src/__tests__/index.test.ts‎

Lines changed: 56 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -80,6 +80,19 @@ vi.mock("../tools/invoices.ts", () => ({
8080
handleListInvoicesRequest: vi.fn(),
8181
handleGetInvoiceRequest: vi.fn(),
8282
}));
83+
vi.mock("../tools/allocations.js", () => ({
84+
listAllocationsTool: {
85+
name: "list_allocations",
86+
description:
87+
"List allocations that your account has access to from the DoiT API",
88+
},
89+
getAllocationTool: {
90+
name: "get_allocation",
91+
description: "Get a specific allocation by ID from the DoiT API",
92+
},
93+
handleListAllocationsRequest: vi.fn(),
94+
handleGetAllocationRequest: vi.fn(),
95+
}));
8396
vi.mock("../utils/util.js", async () => {
8497
const actual = await vi.importActual("../utils/util.js");
8598
return {
@@ -224,6 +237,15 @@ describe("ListToolsRequestSchema Handler", () => {
224237
description:
225238
"Retrieve the full details of an invoice specified by the invoice number from the DoiT API.",
226239
},
240+
{
241+
name: "list_allocations",
242+
description:
243+
"List allocations that your account has access to from the DoiT API",
244+
},
245+
{
246+
name: "get_allocation",
247+
description: "Get a specific allocation by ID from the DoiT API",
248+
},
227249
],
228250
});
229251
});
@@ -490,6 +512,33 @@ describe("CallToolRequestSchema Handler", () => {
490512
expect(handleGetInvoiceRequest).toHaveBeenCalledWith(args, "fake-token");
491513
});
492514

515+
it("should route to the correct tool handler for list_allocations", async () => {
516+
const callToolHandler = setRequestHandlerMock.mock.calls.find(
517+
(call) => call[0] === CallToolRequestSchema
518+
)?.[1];
519+
const args = { pageToken: "next-page-token" };
520+
const request = mockRequest("list_allocations", args);
521+
522+
await callToolHandler(request);
523+
524+
expect(handleListAllocationsRequest).toHaveBeenCalledWith(
525+
args,
526+
"fake-token"
527+
);
528+
});
529+
530+
it("should route to the correct tool handler for get_allocation", async () => {
531+
const callToolHandler = setRequestHandlerMock.mock.calls.find(
532+
(call) => call[0] === CallToolRequestSchema
533+
)?.[1];
534+
const args = { id: "allocation-123" };
535+
const request = mockRequest("get_allocation", args);
536+
537+
await callToolHandler(request);
538+
539+
expect(handleGetAllocationRequest).toHaveBeenCalledWith(args, "fake-token");
540+
});
541+
493542
it("should return Unknown tool error for unknown tool names", async () => {
494543
const callToolHandler = setRequestHandlerMock.mock.calls.find(
495544
(call) => call[0] === CallToolRequestSchema
@@ -589,5 +638,10 @@ describe("InitializeRequestSchema Handler", () => {
589638
});
590639

591640
const indexModule = await import("../index.js");
592-
const { createServer, handleListInvoicesRequest, handleGetInvoiceRequest } =
593-
indexModule;
641+
const {
642+
createServer,
643+
handleListInvoicesRequest,
644+
handleGetInvoiceRequest,
645+
handleListAllocationsRequest,
646+
handleGetAllocationRequest,
647+
} = indexModule;

‎src/index.ts‎

Lines changed: 10 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -57,6 +57,12 @@ import {
5757
getInvoiceTool,
5858
handleGetInvoiceRequest,
5959
} from "./tools/invoices.js";
60+
import {
61+
listAllocationsTool,
62+
handleListAllocationsRequest,
63+
getAllocationTool,
64+
handleGetAllocationRequest,
65+
} from "./tools/allocations.js";
6066

6167
dotenv.config();
6268

@@ -93,6 +99,8 @@ function createServer() {
9399
createTicketTool,
94100
listInvoicesTool,
95101
getInvoiceTool,
102+
listAllocationsTool,
103+
getAllocationTool,
96104
],
97105
};
98106
});
@@ -184,4 +192,6 @@ export {
184192
handleGeneralError,
185193
handleListInvoicesRequest,
186194
handleGetInvoiceRequest,
195+
handleListAllocationsRequest,
196+
handleGetAllocationRequest,
187197
};

‎src/tools/__tests__/allocations.test.ts‎

Lines changed: 231 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,231 @@
1+
import { describe, expect, it, vi, beforeEach } from "vitest";
2+
import {
3+
handleListAllocationsRequest,
4+
handleGetAllocationRequest,
5+
} from "../allocations.js";
6+
import {
7+
createErrorResponse,
8+
createSuccessResponse,
9+
formatZodError,
10+
handleGeneralError,
11+
makeDoitRequest,
12+
} from "../../utils/util.js";
13+
14+
// Mock the utility functions
15+
vi.mock("../../utils/util.js", () => ({
16+
createErrorResponse: vi.fn((msg) => ({
17+
content: [{ type: "text", text: msg }],
18+
})),
19+
createSuccessResponse: vi.fn((text) => ({
20+
content: [{ type: "text", text }],
21+
})),
22+
formatZodError: vi.fn((error) => `Formatted Zod Error: ${error.message}`),
23+
handleGeneralError: vi.fn((error, context) => ({
24+
content: [{ type: "text", text: `General Error: ${context}` }],
25+
})),
26+
makeDoitRequest: vi.fn(),
27+
DOIT_API_BASE: "https://api.doit.com",
28+
}));
29+
30+
describe("allocations", () => {
31+
const mockToken = "fake-token";
32+
33+
beforeEach(() => {
34+
vi.clearAllMocks();
35+
});
36+
37+
describe("handleListAllocationsRequest", () => {
38+
it("should call makeDoitRequest with correct parameters and return success response", async () => {
39+
const mockArgs = { pageToken: "next-page" };
40+
const mockApiResponse = {
41+
pageToken: "another-page",
42+
allocations: [
43+
{
44+
id: "allocation-1",
45+
name: "Test Allocation",
46+
owner: "[email protected]",
47+
type: "standard",
48+
allocationType: "single",
49+
createTime: 1678886400000,
50+
updateTime: 1678972800000,
51+
urlUI: "https://console.doit.com/allocations/allocation-1",
52+
},
53+
],
54+
};
55+
(makeDoitRequest as vi.Mock).mockResolvedValue(mockApiResponse);
56+
57+
const response = await handleListAllocationsRequest(mockArgs, mockToken);
58+
59+
expect(makeDoitRequest).toHaveBeenCalledWith(
60+
"https://api.doit.com/analytics/v1/allocations?pageToken=next-page",
61+
mockToken,
62+
{
63+
method: "GET",
64+
}
65+
);
66+
expect(createSuccessResponse).toHaveBeenCalledWith(
67+
expect.stringContaining("allocation-1")
68+
);
69+
});
70+
71+
it("should handle request without pageToken", async () => {
72+
const mockArgs = {};
73+
const mockApiResponse = {
74+
allocations: [
75+
{
76+
id: "allocation-1",
77+
name: "Test Allocation",
78+
owner: "[email protected]",
79+
type: "standard",
80+
allocationType: "single",
81+
createTime: 1678886400000,
82+
updateTime: 1678972800000,
83+
urlUI: "https://console.doit.com/allocations/allocation-1",
84+
},
85+
],
86+
};
87+
(makeDoitRequest as vi.Mock).mockResolvedValue(mockApiResponse);
88+
89+
const response = await handleListAllocationsRequest(mockArgs, mockToken);
90+
91+
expect(makeDoitRequest).toHaveBeenCalledWith(
92+
"https://api.doit.com/analytics/v1/allocations",
93+
mockToken,
94+
{
95+
method: "GET",
96+
}
97+
);
98+
expect(createSuccessResponse).toHaveBeenCalled();
99+
});
100+
101+
it("should handle no allocations found", async () => {
102+
const mockArgs = {};
103+
const mockApiResponse = {
104+
allocations: [],
105+
};
106+
(makeDoitRequest as vi.Mock).mockResolvedValue(mockApiResponse);
107+
108+
const response = await handleListAllocationsRequest(mockArgs, mockToken);
109+
110+
expect(createErrorResponse).toHaveBeenCalledWith("No allocations found");
111+
});
112+
113+
it("should handle API request failure", async () => {
114+
const mockArgs = {};
115+
(makeDoitRequest as vi.Mock).mockResolvedValue(null);
116+
117+
const response = await handleListAllocationsRequest(mockArgs, mockToken);
118+
119+
expect(createErrorResponse).toHaveBeenCalledWith(
120+
"Failed to retrieve allocations data"
121+
);
122+
});
123+
124+
it("should handle makeDoitRequest throwing an error", async () => {
125+
const mockArgs = {};
126+
(makeDoitRequest as vi.Mock).mockRejectedValue(
127+
new Error("Network error")
128+
);
129+
130+
const response = await handleListAllocationsRequest(mockArgs, mockToken);
131+
132+
expect(handleGeneralError).toHaveBeenCalledWith(
133+
expect.any(Error),
134+
"making DoiT API request"
135+
);
136+
});
137+
});
138+
139+
describe("handleGetAllocationRequest", () => {
140+
it("should call makeDoitRequest with correct parameters and return success response", async () => {
141+
const mockArgs = { id: "allocation-123" };
142+
const mockApiResponse = {
143+
id: "allocation-123",
144+
name: "Test Allocation",
145+
description: "A test allocation",
146+
type: "standard",
147+
allocationType: "single",
148+
createTime: 1678886400000,
149+
updateTime: 1678972800000,
150+
anomalyDetection: true,
151+
rule: {
152+
components: [
153+
{
154+
key: "project_id",
155+
type: "fixed",
156+
values: ["project-1", "project-2"],
157+
inverse_selection: false,
158+
include_null: false,
159+
mode: "is",
160+
},
161+
],
162+
formula: "A AND B",
163+
},
164+
};
165+
(makeDoitRequest as vi.Mock).mockResolvedValue(mockApiResponse);
166+
167+
const response = await handleGetAllocationRequest(mockArgs, mockToken);
168+
169+
expect(makeDoitRequest).toHaveBeenCalledWith(
170+
"https://api.doit.com/analytics/v1/allocations/allocation-123",
171+
mockToken,
172+
{
173+
method: "GET",
174+
}
175+
);
176+
expect(createSuccessResponse).toHaveBeenCalledWith(
177+
expect.stringContaining("allocation-123")
178+
);
179+
});
180+
181+
it("should handle missing allocation ID", async () => {
182+
const mockArgs = {};
183+
184+
const response = await handleGetAllocationRequest(mockArgs, mockToken);
185+
186+
expect(createErrorResponse).toHaveBeenCalledWith(
187+
expect.stringContaining("Formatted Zod Error")
188+
);
189+
});
190+
191+
it("should handle API request failure", async () => {
192+
const mockArgs = { id: "allocation-123" };
193+
(makeDoitRequest as vi.Mock).mockResolvedValue(null);
194+
195+
const response = await handleGetAllocationRequest(mockArgs, mockToken);
196+
197+
expect(createErrorResponse).toHaveBeenCalledWith(
198+
"Failed to retrieve allocation data"
199+
);
200+
});
201+
202+
it("should handle makeDoitRequest throwing an error", async () => {
203+
const mockArgs = { id: "allocation-123" };
204+
(makeDoitRequest as vi.Mock).mockRejectedValue(
205+
new Error("Network error")
206+
);
207+
208+
const response = await handleGetAllocationRequest(mockArgs, mockToken);
209+
210+
expect(handleGeneralError).toHaveBeenCalledWith(
211+
expect.any(Error),
212+
"making DoiT API request"
213+
);
214+
});
215+
216+
it("should handle Zod validation errors", async () => {
217+
const mockArgs = { id: 123 }; // Invalid type (should be string)
218+
219+
const response = await handleGetAllocationRequest(mockArgs, mockToken);
220+
221+
expect(response).toEqual({
222+
content: [
223+
{
224+
type: "text",
225+
text: expect.stringContaining("Formatted Zod Error"),
226+
},
227+
],
228+
});
229+
});
230+
});
231+
});

0 commit comments

Comments
 (0)