Skip to content

Commit

Permalink
Merge pull request #1840 from ai16z-demirix/test/goals-memory-provider
Browse files Browse the repository at this point in the history 
test: adding more tests for goals,memory and provider. Fixing generation.test.ts
  • Loading branch information
@shakkernerd
shakkernerd authored Jan 5, 2025
2 parents 9f4694f + 5e78b7a commit ed16b0c
Show file tree
Hide file tree
Showing 4 changed files with 477 additions and 5 deletions.
25 changes: 21 additions & 4 deletions packages/core/src/tests/generation.test.ts
View file
Original file line number Diff line number Diff line change
Expand Up @@ -129,6 +129,23 @@ describe("Generation", () => {
});

describe("trimTokens", () => {
let mockRuntime: IAgentRuntime;

beforeEach(() => {
mockRuntime = {
getSetting: vi.fn().mockImplementation((key: string) => {
switch (key) {
case "TOKENIZER_MODEL":
return "gpt-4";
case "TOKENIZER_TYPE":
return "tiktoken";
default:
return undefined;
}
}),
} as unknown as IAgentRuntime;
});

it("should return empty string for empty input", async () => {
const result = await trimTokens("", 100, mockRuntime);
expect(result).toBe("");
Expand Down Expand Up @@ -169,10 +186,10 @@ describe("Generation", () => {

it("should handle multiline text", async () => {
const multilineText = `Line 1
Line 2
Line 3
Line 4
Line 5`;
Line 2
Line 3
Line 4
Line 5`;
const result = await trimTokens(multilineText, 5, mockRuntime);
expect(result.length).toBeGreaterThan(0);
expect(result.length).toBeLessThan(multilineText.length);
Expand Down
261 changes: 260 additions & 1 deletion packages/core/src/tests/goals.test.ts
View file
Original file line number Diff line number Diff line change
Expand Up @@ -247,8 +247,15 @@ const sampleGoal: Goal = {
};

describe("getGoals", () => {
let runtime: IAgentRuntime;

beforeEach(() => {
vi.clearAllMocks();
runtime = {
agentId: "test-agent-id" as UUID,
databaseAdapter: {
getGoals: vi.fn().mockResolvedValue([]),
} as any,
} as IAgentRuntime;
});

it("retrieves goals successfully", async () => {
Expand All @@ -274,6 +281,26 @@ describe("getGoals", () => {
})
).rejects.toThrow("Failed to retrieve goals");
});

it("should handle empty goals list", async () => {
const mockRuntime = {
agentId: "test-agent-id" as UUID,
databaseAdapter: {
getGoals: vi.fn().mockResolvedValue([]),
},
} as unknown as IAgentRuntime;

const roomId = "test-room" as UUID;

await getGoals({ runtime: mockRuntime, roomId });

expect(mockRuntime.databaseAdapter.getGoals).toHaveBeenCalledWith({
agentId: "test-agent-id",
roomId,
onlyInProgress: true,
count: 5,
});
});
});

describe("formatGoalsAsString", () => {
Expand All @@ -292,6 +319,53 @@ describe("formatGoalsAsString", () => {
const formatted = formatGoalsAsString({ goals: [] });
expect(formatted).toBe("");
});

it("should format goals as string correctly", () => {
const goals: Goal[] = [
{
id: "1" as UUID,
name: "Goal 1",
status: GoalStatus.IN_PROGRESS,
objectives: [
{
id: "obj1" as UUID,
description: "Objective 1",
completed: true,
},
{
id: "obj2" as UUID,
description: "Objective 2",
completed: false,
},
],
roomId: "test-room" as UUID,
userId: "test-user" as UUID,
},
{
id: "2" as UUID,
name: "Goal 2",
status: GoalStatus.DONE,
objectives: [
{
id: "obj3" as UUID,
description: "Objective 3",
completed: true,
},
],
roomId: "test-room" as UUID,
userId: "test-user" as UUID,
},
];

const formattedGoals = formatGoalsAsString({ goals });
expect(formattedGoals).toContain("Goal: Goal 1");
expect(formattedGoals).toContain("id: 1");
expect(formattedGoals).toContain("- [x] Objective 1 (DONE)");
expect(formattedGoals).toContain("- [ ] Objective 2 (IN PROGRESS)");
expect(formattedGoals).toContain("Goal: Goal 2");
expect(formattedGoals).toContain("id: 2");
expect(formattedGoals).toContain("- [x] Objective 3 (DONE)");
});
});

describe("updateGoal", () => {
Expand All @@ -318,6 +392,138 @@ describe("updateGoal", () => {
updateGoal({ runtime: mockRuntime, goal: sampleGoal })
).rejects.toThrow("Failed to update goal");
});

it("should update goal status correctly", async () => {
const goalId = "test-goal" as UUID;
const mockRuntime = {
databaseAdapter: { updateGoal: vi.fn() },
agentId: "test-agent-id" as UUID,
} as unknown as IAgentRuntime;

const updatedGoal: Goal = {
id: goalId,
name: "Test Goal",
objectives: [
{
description: "Objective 1",
completed: false,
},
{
description: "Objective 2",
completed: true,
},
],
roomId: "room-id" as UUID,
userId: "user-id" as UUID,
status: GoalStatus.DONE,
};

await updateGoal({
runtime: mockRuntime,
goal: updatedGoal,
});

expect(mockRuntime.databaseAdapter.updateGoal).toHaveBeenCalledWith(updatedGoal);
});

it("should handle failed goal update", async () => {
const goalId = "test-goal" as UUID;
const mockRuntime = {
databaseAdapter: { updateGoal: vi.fn() },
agentId: "test-agent-id" as UUID,
} as unknown as IAgentRuntime;

const updatedGoal: Goal = {
id: goalId,
name: "Test Goal",
objectives: [
{
description: "Objective 1",
completed: false,
},
{
description: "Objective 2",
completed: true,
},
],
roomId: "room-id" as UUID,
userId: "user-id" as UUID,
status: GoalStatus.FAILED,
};

await updateGoal({
runtime: mockRuntime,
goal: updatedGoal,
});

expect(mockRuntime.databaseAdapter.updateGoal).toHaveBeenCalledWith(updatedGoal);
});

it("should handle in-progress goal update", async () => {
const goalId = "test-goal" as UUID;
const mockRuntime = {
databaseAdapter: { updateGoal: vi.fn() },
agentId: "test-agent-id" as UUID,
} as unknown as IAgentRuntime;

const updatedGoal: Goal = {
id: goalId,
name: "Test Goal",
objectives: [
{
description: "Objective 1",
completed: false,
},
{
description: "Objective 2",
completed: true,
},
],
roomId: "room-id" as UUID,
userId: "user-id" as UUID,
status: GoalStatus.IN_PROGRESS,
};

await updateGoal({
runtime: mockRuntime,
goal: updatedGoal,
});

expect(mockRuntime.databaseAdapter.updateGoal).toHaveBeenCalledWith(updatedGoal);
});

it("should handle goal priority updates", async () => {
const goalId = "test-goal" as UUID;
const mockRuntime = {
databaseAdapter: { updateGoal: vi.fn() },
agentId: "test-agent-id" as UUID,
} as unknown as IAgentRuntime;

const updatedGoal: Goal = {
id: goalId,
name: "Test Goal",
objectives: [
{
description: "Objective 1",
completed: false,
},
{
description: "Objective 2",
completed: true,
},
],
roomId: "room-id" as UUID,
userId: "user-id" as UUID,
status: GoalStatus.IN_PROGRESS,
};

await updateGoal({
runtime: mockRuntime,
goal: updatedGoal,
});

expect(mockRuntime.databaseAdapter.updateGoal).toHaveBeenCalledWith(updatedGoal);
});
});

describe("createGoal", () => {
Expand All @@ -344,4 +550,57 @@ describe("createGoal", () => {
createGoal({ runtime: mockRuntime, goal: sampleGoal })
).rejects.toThrow("Failed to create goal");
});

it("should create new goal with correct properties", async () => {
const newGoal: Goal = {
name: "New Goal",
roomId: "room-id" as UUID,
userId: "user-id" as UUID,
status: GoalStatus.IN_PROGRESS,
objectives: []
};

const mockRuntime = {
databaseAdapter: { createGoal: vi.fn() },
agentId: "test-agent-id" as UUID,
} as unknown as IAgentRuntime;

await createGoal({
runtime: mockRuntime,
goal: newGoal,
});

expect(mockRuntime.databaseAdapter.createGoal).toHaveBeenCalledWith(
expect.objectContaining({
name: "New Goal",
roomId: "room-id",
userId: "user-id",
status: GoalStatus.IN_PROGRESS,
objectives: []
})
);
});

it("should create a new goal", async () => {
const mockRuntime = {
databaseAdapter: { createGoal: vi.fn() },
agentId: "test-agent-id" as UUID,
} as unknown as IAgentRuntime;

const newGoal = {
id: "new-goal" as UUID,
name: "New Goal",
objectives: [],
roomId: "test-room" as UUID,
userId: "test-user" as UUID,
status: GoalStatus.IN_PROGRESS,
};

await createGoal({
runtime: mockRuntime,
goal: newGoal,
});

expect(mockRuntime.databaseAdapter.createGoal).toHaveBeenCalledWith(newGoal);
});
});
Loading

0 comments on commit ed16b0c

Please sign in to comment.