stack/apps/api/src/agent-memory/agent-memory.controller.spec.ts

import { Test, TestingModule } from "@nestjs/testing";
import { AgentMemoryController } from "./agent-memory.controller";
import { AgentMemoryService } from "./agent-memory.service";
import { AuthGuard } from "../auth/guards/auth.guard";
import { WorkspaceGuard, PermissionGuard } from "../common/guards";
import { describe, it, expect, beforeEach, vi } from "vitest";

describe("AgentMemoryController", () => {
  let controller: AgentMemoryController;

  const mockAgentMemoryService = {
    upsert: vi.fn(),
    findAll: vi.fn(),
    findOne: vi.fn(),
    remove: vi.fn(),
  };

  const mockGuard = { canActivate: vi.fn(() => true) };

  beforeEach(async () => {
    const module: TestingModule = await Test.createTestingModule({
      controllers: [AgentMemoryController],
      providers: [
        {
          provide: AgentMemoryService,
          useValue: mockAgentMemoryService,
        },
      ],
    })
      .overrideGuard(AuthGuard)
      .useValue(mockGuard)
      .overrideGuard(WorkspaceGuard)
      .useValue(mockGuard)
      .overrideGuard(PermissionGuard)
      .useValue(mockGuard)
      .compile();

    controller = module.get<AgentMemoryController>(AgentMemoryController);

    vi.clearAllMocks();
  });

  const workspaceId = "workspace-1";
  const agentId = "agent-1";
  const key = "context";

  describe("upsert", () => {
    it("should upsert a memory entry", async () => {
      const dto = { value: { foo: "bar" } };
      const mockEntry = { id: "mem-1", workspaceId, agentId, key, value: dto.value };

      mockAgentMemoryService.upsert.mockResolvedValue(mockEntry);

      const result = await controller.upsert(agentId, key, dto, workspaceId);

      expect(mockAgentMemoryService.upsert).toHaveBeenCalledWith(workspaceId, agentId, key, dto);
      expect(result).toEqual(mockEntry);
    });
  });

  describe("findAll", () => {
    it("should list all memory entries for an agent", async () => {
      const mockEntries = [
        { id: "mem-1", key: "a", value: 1 },
        { id: "mem-2", key: "b", value: 2 },
      ];

      mockAgentMemoryService.findAll.mockResolvedValue(mockEntries);

      const result = await controller.findAll(agentId, workspaceId);

      expect(mockAgentMemoryService.findAll).toHaveBeenCalledWith(workspaceId, agentId);
      expect(result).toEqual(mockEntries);
    });
  });

  describe("findOne", () => {
    it("should get a single memory entry", async () => {
      const mockEntry = { id: "mem-1", key, value: "v" };

      mockAgentMemoryService.findOne.mockResolvedValue(mockEntry);

      const result = await controller.findOne(agentId, key, workspaceId);

      expect(mockAgentMemoryService.findOne).toHaveBeenCalledWith(workspaceId, agentId, key);
      expect(result).toEqual(mockEntry);
    });
  });

  describe("remove", () => {
    it("should delete a memory entry", async () => {
      const mockResponse = { message: "Memory entry deleted successfully" };

      mockAgentMemoryService.remove.mockResolvedValue(mockResponse);

      const result = await controller.remove(agentId, key, workspaceId);

      expect(mockAgentMemoryService.remove).toHaveBeenCalledWith(workspaceId, agentId, key);
      expect(result).toEqual(mockResponse);
    });
  });
});