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

import { Test, TestingModule } from "@nestjs/testing";
import { AgentMemoryService } from "./agent-memory.service";
import { PrismaService } from "../prisma/prisma.service";
import { NotFoundException } from "@nestjs/common";
import { describe, it, expect, beforeEach, vi } from "vitest";

describe("AgentMemoryService", () => {
  let service: AgentMemoryService;

  const mockPrismaService = {
    agentMemory: {
      upsert: vi.fn(),
      findMany: vi.fn(),
      findUnique: vi.fn(),
      delete: vi.fn(),
    },
  };

  beforeEach(async () => {
    const module: TestingModule = await Test.createTestingModule({
      providers: [
        AgentMemoryService,
        {
          provide: PrismaService,
          useValue: mockPrismaService,
        },
      ],
    }).compile();

    service = module.get<AgentMemoryService>(AgentMemoryService);

    vi.clearAllMocks();
  });

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

  describe("upsert", () => {
    it("should upsert a memory entry", async () => {
      const dto = { value: { data: "some context" } };
      const mockEntry = {
        id: "mem-1",
        workspaceId,
        agentId,
        key,
        value: dto.value,
        createdAt: new Date(),
        updatedAt: new Date(),
      };

      mockPrismaService.agentMemory.upsert.mockResolvedValue(mockEntry);

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

      expect(mockPrismaService.agentMemory.upsert).toHaveBeenCalledWith({
        where: { workspaceId_agentId_key: { workspaceId, agentId, key } },
        create: { workspaceId, agentId, key, value: dto.value },
        update: { value: dto.value },
      });
      expect(result).toEqual(mockEntry);
    });
  });

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

      mockPrismaService.agentMemory.findMany.mockResolvedValue(mockEntries);

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

      expect(mockPrismaService.agentMemory.findMany).toHaveBeenCalledWith({
        where: { workspaceId, agentId },
        orderBy: { key: "asc" },
      });
      expect(result).toEqual(mockEntries);
    });
  });

  describe("findOne", () => {
    it("should return a memory entry by key", async () => {
      const mockEntry = { id: "mem-1", workspaceId, agentId, key, value: "ctx" };

      mockPrismaService.agentMemory.findUnique.mockResolvedValue(mockEntry);

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

      expect(mockPrismaService.agentMemory.findUnique).toHaveBeenCalledWith({
        where: { workspaceId_agentId_key: { workspaceId, agentId, key } },
      });
      expect(result).toEqual(mockEntry);
    });

    it("should throw NotFoundException when key not found", async () => {
      mockPrismaService.agentMemory.findUnique.mockResolvedValue(null);

      await expect(service.findOne(workspaceId, agentId, key)).rejects.toThrow(NotFoundException);
    });
  });

  describe("remove", () => {
    it("should delete a memory entry", async () => {
      const mockEntry = { id: "mem-1", workspaceId, agentId, key, value: "x" };

      mockPrismaService.agentMemory.findUnique.mockResolvedValue(mockEntry);
      mockPrismaService.agentMemory.delete.mockResolvedValue(mockEntry);

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

      expect(mockPrismaService.agentMemory.delete).toHaveBeenCalledWith({
        where: { workspaceId_agentId_key: { workspaceId, agentId, key } },
      });
      expect(result).toEqual({ message: "Memory entry deleted successfully" });
    });

    it("should throw NotFoundException when key not found", async () => {
      mockPrismaService.agentMemory.findUnique.mockResolvedValue(null);

      await expect(service.remove(workspaceId, agentId, key)).rejects.toThrow(NotFoundException);
    });
  });
});