stack/apps/api/src/brain/brain.controller.test.ts

import { describe, expect, it, vi, beforeEach } from "vitest";
import { BrainController } from "./brain.controller";
import { BrainService, BrainQueryResult, BrainContext } from "./brain.service";
import { IntentClassificationService } from "./intent-classification.service";
import type { IntentClassification } from "./interfaces";
import { TaskStatus, TaskPriority, ProjectStatus, EntityType } from "@prisma/client";

describe("BrainController", () => {
  let controller: BrainController;
  let mockService: {
    query: ReturnType<typeof vi.fn>;
    getContext: ReturnType<typeof vi.fn>;
    search: ReturnType<typeof vi.fn>;
  };
  let mockIntentService: {
    classify: ReturnType<typeof vi.fn>;
  };

  const mockWorkspaceId = "123e4567-e89b-12d3-a456-426614174000";

  const mockQueryResult: BrainQueryResult = {
    tasks: [
      {
        id: "task-1",
        title: "Test Task",
        description: null,
        status: TaskStatus.IN_PROGRESS,
        priority: TaskPriority.HIGH,
        dueDate: null,
        assignee: null,
        project: null,
      },
    ],
    events: [
      {
        id: "event-1",
        title: "Test Event",
        description: null,
        startTime: new Date("2025-02-01T10:00:00Z"),
        endTime: new Date("2025-02-01T11:00:00Z"),
        allDay: false,
        location: null,
        project: null,
      },
    ],
    projects: [
      {
        id: "project-1",
        name: "Test Project",
        description: null,
        status: ProjectStatus.ACTIVE,
        startDate: null,
        endDate: null,
        color: null,
        _count: { tasks: 5, events: 2 },
      },
    ],
    meta: {
      totalTasks: 1,
      totalEvents: 1,
      totalProjects: 1,
      filters: {},
    },
  };

  const mockContext: BrainContext = {
    timestamp: new Date(),
    workspace: { id: mockWorkspaceId, name: "Test Workspace" },
    summary: {
      activeTasks: 10,
      overdueTasks: 2,
      upcomingEvents: 5,
      activeProjects: 3,
    },
    tasks: [
      {
        id: "task-1",
        title: "Test Task",
        status: TaskStatus.IN_PROGRESS,
        priority: TaskPriority.HIGH,
        dueDate: null,
        isOverdue: false,
      },
    ],
    events: [
      {
        id: "event-1",
        title: "Test Event",
        startTime: new Date("2025-02-01T10:00:00Z"),
        endTime: new Date("2025-02-01T11:00:00Z"),
        allDay: false,
        location: null,
      },
    ],
    projects: [
      {
        id: "project-1",
        name: "Test Project",
        status: ProjectStatus.ACTIVE,
        taskCount: 5,
      },
    ],
  };

  const mockIntentResult: IntentClassification = {
    intent: "query_tasks",
    confidence: 0.9,
    entities: [],
    method: "rule",
    query: "show my tasks",
  };

  beforeEach(() => {
    mockService = {
      query: vi.fn().mockResolvedValue(mockQueryResult),
      getContext: vi.fn().mockResolvedValue(mockContext),
      search: vi.fn().mockResolvedValue(mockQueryResult),
    };

    mockIntentService = {
      classify: vi.fn().mockResolvedValue(mockIntentResult),
    };

    controller = new BrainController(
      mockService as unknown as BrainService,
      mockIntentService as unknown as IntentClassificationService
    );
  });

  describe("query", () => {
    it("should call service.query with merged workspaceId", async () => {
      const queryDto = {
        workspaceId: "different-id",
        query: "What tasks are due?",
      };

      const result = await controller.query(queryDto, mockWorkspaceId);

      expect(mockService.query).toHaveBeenCalledWith({
        ...queryDto,
        workspaceId: mockWorkspaceId,
      });
      expect(result).toEqual(mockQueryResult);
    });

    it("should handle query with filters", async () => {
      const queryDto = {
        workspaceId: mockWorkspaceId,
        entities: [EntityType.TASK, EntityType.EVENT],
        tasks: { status: TaskStatus.IN_PROGRESS },
        events: { upcoming: true },
      };

      await controller.query(queryDto, mockWorkspaceId);

      expect(mockService.query).toHaveBeenCalledWith({
        ...queryDto,
        workspaceId: mockWorkspaceId,
      });
    });

    it("should handle query with search term", async () => {
      const queryDto = {
        workspaceId: mockWorkspaceId,
        search: "important",
        limit: 10,
      };

      await controller.query(queryDto, mockWorkspaceId);

      expect(mockService.query).toHaveBeenCalledWith({
        ...queryDto,
        workspaceId: mockWorkspaceId,
      });
    });

    it("should return query result structure", async () => {
      const result = await controller.query({ workspaceId: mockWorkspaceId }, mockWorkspaceId);

      expect(result).toHaveProperty("tasks");
      expect(result).toHaveProperty("events");
      expect(result).toHaveProperty("projects");
      expect(result).toHaveProperty("meta");
      expect(result.tasks).toHaveLength(1);
      expect(result.events).toHaveLength(1);
      expect(result.projects).toHaveLength(1);
    });
  });

  describe("getContext", () => {
    it("should call service.getContext with merged workspaceId", async () => {
      const contextDto = {
        workspaceId: "different-id",
        includeTasks: true,
      };

      const result = await controller.getContext(contextDto, mockWorkspaceId);

      expect(mockService.getContext).toHaveBeenCalledWith({
        ...contextDto,
        workspaceId: mockWorkspaceId,
      });
      expect(result).toEqual(mockContext);
    });

    it("should handle context with all options", async () => {
      const contextDto = {
        workspaceId: mockWorkspaceId,
        includeTasks: true,
        includeEvents: true,
        includeProjects: true,
        eventDays: 14,
      };

      await controller.getContext(contextDto, mockWorkspaceId);

      expect(mockService.getContext).toHaveBeenCalledWith({
        ...contextDto,
        workspaceId: mockWorkspaceId,
      });
    });

    it("should return context structure", async () => {
      const result = await controller.getContext({ workspaceId: mockWorkspaceId }, mockWorkspaceId);

      expect(result).toHaveProperty("timestamp");
      expect(result).toHaveProperty("workspace");
      expect(result).toHaveProperty("summary");
      expect(result.summary).toHaveProperty("activeTasks");
      expect(result.summary).toHaveProperty("overdueTasks");
      expect(result.summary).toHaveProperty("upcomingEvents");
      expect(result.summary).toHaveProperty("activeProjects");
    });

    it("should include detailed lists when requested", async () => {
      const result = await controller.getContext(
        {
          workspaceId: mockWorkspaceId,
          includeTasks: true,
          includeEvents: true,
          includeProjects: true,
        },
        mockWorkspaceId
      );

      expect(result.tasks).toBeDefined();
      expect(result.events).toBeDefined();
      expect(result.projects).toBeDefined();
    });
  });

  describe("search", () => {
    it("should call service.search with parameters from DTO", async () => {
      const result = await controller.search({ q: "test query", limit: 10 }, mockWorkspaceId);

      expect(mockService.search).toHaveBeenCalledWith(mockWorkspaceId, "test query", 10);
      expect(result).toEqual(mockQueryResult);
    });

    it("should use default limit when not provided in DTO", async () => {
      await controller.search({ q: "test" }, mockWorkspaceId);

      expect(mockService.search).toHaveBeenCalledWith(mockWorkspaceId, "test", 20);
    });

    it("should handle empty search DTO", async () => {
      await controller.search({}, mockWorkspaceId);

      expect(mockService.search).toHaveBeenCalledWith(mockWorkspaceId, "", 20);
    });

    it("should handle undefined q in DTO", async () => {
      await controller.search({ limit: 10 }, mockWorkspaceId);

      expect(mockService.search).toHaveBeenCalledWith(mockWorkspaceId, "", 10);
    });

    it("should return search result structure", async () => {
      const result = await controller.search({ q: "test", limit: 10 }, mockWorkspaceId);

      expect(result).toHaveProperty("tasks");
      expect(result).toHaveProperty("events");
      expect(result).toHaveProperty("projects");
      expect(result).toHaveProperty("meta");
    });
  });

  describe("classifyIntent", () => {
    it("should call intentService.classify with query", async () => {
      const dto = { query: "show my tasks" };

      const result = await controller.classifyIntent(dto);

      expect(mockIntentService.classify).toHaveBeenCalledWith("show my tasks", undefined);
      expect(result).toEqual(mockIntentResult);
    });

    it("should pass useLlm flag when provided", async () => {
      const dto = { query: "show my tasks", useLlm: true };

      await controller.classifyIntent(dto);

      expect(mockIntentService.classify).toHaveBeenCalledWith("show my tasks", true);
    });

    it("should return intent classification structure", async () => {
      const result = await controller.classifyIntent({ query: "show my tasks" });

      expect(result).toHaveProperty("intent");
      expect(result).toHaveProperty("confidence");
      expect(result).toHaveProperty("entities");
      expect(result).toHaveProperty("method");
      expect(result).toHaveProperty("query");
    });

    it("should handle different intent types", async () => {
      const briefingResult: IntentClassification = {
        intent: "briefing",
        confidence: 0.95,
        entities: [],
        method: "rule",
        query: "morning briefing",
      };
      mockIntentService.classify.mockResolvedValue(briefingResult);

      const result = await controller.classifyIntent({ query: "morning briefing" });

      expect(result.intent).toBe("briefing");
      expect(result.confidence).toBe(0.95);
    });

    it("should handle intent with entities", async () => {
      const resultWithEntities: IntentClassification = {
        intent: "create_task",
        confidence: 0.9,
        entities: [
          {
            type: "priority",
            value: "HIGH",
            raw: "high priority",
            start: 12,
            end: 25,
          },
        ],
        method: "rule",
        query: "create task high priority",
      };
      mockIntentService.classify.mockResolvedValue(resultWithEntities);

      const result = await controller.classifyIntent({ query: "create task high priority" });

      expect(result.entities).toHaveLength(1);
      expect(result.entities[0].type).toBe("priority");
      expect(result.entities[0].value).toBe("HIGH");
    });

    it("should handle LLM classification", async () => {
      const llmResult: IntentClassification = {
        intent: "search",
        confidence: 0.85,
        entities: [],
        method: "llm",
        query: "find something",
      };
      mockIntentService.classify.mockResolvedValue(llmResult);

      const result = await controller.classifyIntent({ query: "find something", useLlm: true });

      expect(result.method).toBe("llm");
      expect(result.intent).toBe("search");
    });
  });
});