stack/apps/orchestrator/src/api/agents/agents.controller.spec.ts

import { AgentsController } from "./agents.controller";
import { QueueService } from "../../queue/queue.service";
import { AgentSpawnerService } from "../../spawner/agent-spawner.service";
import { AgentLifecycleService } from "../../spawner/agent-lifecycle.service";
import { KillswitchService } from "../../killswitch/killswitch.service";
import { BadRequestException } from "@nestjs/common";
import { describe, it, expect, beforeEach, afterEach, vi } from "vitest";

describe("AgentsController", () => {
  let controller: AgentsController;
  let queueService: {
    addTask: ReturnType<typeof vi.fn>;
  };
  let spawnerService: {
    spawnAgent: ReturnType<typeof vi.fn>;
    listAgentSessions: ReturnType<typeof vi.fn>;
    getAgentSession: ReturnType<typeof vi.fn>;
  };
  let lifecycleService: {
    getAgentLifecycleState: ReturnType<typeof vi.fn>;
  };
  let killswitchService: {
    killAgent: ReturnType<typeof vi.fn>;
    killAllAgents: ReturnType<typeof vi.fn>;
  };

  beforeEach(() => {
    // Create mock services
    queueService = {
      addTask: vi.fn().mockResolvedValue(undefined),
    };

    spawnerService = {
      spawnAgent: vi.fn(),
      listAgentSessions: vi.fn(),
      getAgentSession: vi.fn(),
    };

    lifecycleService = {
      getAgentLifecycleState: vi.fn(),
    };

    killswitchService = {
      killAgent: vi.fn(),
      killAllAgents: vi.fn(),
    };

    // Create controller with mocked services
    controller = new AgentsController(
      queueService as unknown as QueueService,
      spawnerService as unknown as AgentSpawnerService,
      lifecycleService as unknown as AgentLifecycleService,
      killswitchService as unknown as KillswitchService
    );
  });

  afterEach(() => {
    vi.clearAllMocks();
  });

  it("should be defined", () => {
    expect(controller).toBeDefined();
  });

  describe("listAgents", () => {
    it("should return empty array when no agents exist", () => {
      // Arrange
      spawnerService.listAgentSessions.mockReturnValue([]);

      // Act
      const result = controller.listAgents();

      // Assert
      expect(spawnerService.listAgentSessions).toHaveBeenCalled();
      expect(result).toEqual([]);
    });

    it("should return all agent sessions with mapped status", () => {
      // Arrange
      const sessions = [
        {
          agentId: "agent-1",
          taskId: "task-1",
          agentType: "worker" as const,
          state: "running" as const,
          context: {
            repository: "repo",
            branch: "main",
            workItems: [],
          },
          spawnedAt: new Date("2026-02-05T12:00:00Z"),
        },
        {
          agentId: "agent-2",
          taskId: "task-2",
          agentType: "reviewer" as const,
          state: "completed" as const,
          context: {
            repository: "repo",
            branch: "main",
            workItems: [],
          },
          spawnedAt: new Date("2026-02-05T11:00:00Z"),
          completedAt: new Date("2026-02-05T11:30:00Z"),
        },
        {
          agentId: "agent-3",
          taskId: "task-3",
          agentType: "tester" as const,
          state: "failed" as const,
          context: {
            repository: "repo",
            branch: "main",
            workItems: [],
          },
          spawnedAt: new Date("2026-02-05T10:00:00Z"),
          error: "Test execution failed",
        },
      ];
      spawnerService.listAgentSessions.mockReturnValue(sessions);

      // Act
      const result = controller.listAgents();

      // Assert
      expect(spawnerService.listAgentSessions).toHaveBeenCalled();
      expect(result).toHaveLength(3);
      expect(result[0]).toEqual({
        agentId: "agent-1",
        taskId: "task-1",
        status: "running",
        agentType: "worker",
        spawnedAt: "2026-02-05T12:00:00.000Z",
        completedAt: undefined,
        error: undefined,
      });
      expect(result[1]).toEqual({
        agentId: "agent-2",
        taskId: "task-2",
        status: "completed",
        agentType: "reviewer",
        spawnedAt: "2026-02-05T11:00:00.000Z",
        completedAt: "2026-02-05T11:30:00.000Z",
        error: undefined,
      });
      expect(result[2]).toEqual({
        agentId: "agent-3",
        taskId: "task-3",
        status: "failed",
        agentType: "tester",
        spawnedAt: "2026-02-05T10:00:00.000Z",
        completedAt: undefined,
        error: "Test execution failed",
      });
    });

    it("should handle errors gracefully", () => {
      // Arrange
      spawnerService.listAgentSessions.mockImplementation(() => {
        throw new Error("Service unavailable");
      });

      // Act & Assert
      expect(() => controller.listAgents()).toThrow("Failed to list agents: Service unavailable");
    });
  });

  describe("spawn", () => {
    const validRequest = {
      taskId: "task-123",
      agentType: "worker" as const,
      context: {
        repository: "https://github.com/org/repo.git",
        branch: "main",
        workItems: ["US-001", "US-002"],
        skills: ["typescript", "nestjs"],
      },
    };

    it("should spawn agent and queue task successfully", async () => {
      // Arrange
      const agentId = "agent-abc-123";
      const spawnedAt = new Date();
      spawnerService.spawnAgent.mockReturnValue({
        agentId,
        state: "spawning",
        spawnedAt,
      });
      queueService.addTask.mockResolvedValue(undefined);

      // Act
      const result = await controller.spawn(validRequest);

      // Assert
      expect(spawnerService.spawnAgent).toHaveBeenCalledWith(validRequest);
      expect(queueService.addTask).toHaveBeenCalledWith(validRequest.taskId, validRequest.context, {
        priority: 5,
      });
      expect(result).toEqual({
        agentId,
        status: "spawning",
      });
    });

    it("should return queued status when agent is queued", async () => {
      // Arrange
      const agentId = "agent-abc-123";
      spawnerService.spawnAgent.mockReturnValue({
        agentId,
        state: "spawning",
        spawnedAt: new Date(),
      });
      queueService.addTask.mockResolvedValue(undefined);

      // Act
      const result = await controller.spawn(validRequest);

      // Assert
      expect(result.status).toBe("spawning");
    });

    it("should handle reviewer agent type", async () => {
      // Arrange
      const reviewerRequest = {
        ...validRequest,
        agentType: "reviewer" as const,
      };
      const agentId = "agent-reviewer-123";
      spawnerService.spawnAgent.mockReturnValue({
        agentId,
        state: "spawning",
        spawnedAt: new Date(),
      });
      queueService.addTask.mockResolvedValue(undefined);

      // Act
      const result = await controller.spawn(reviewerRequest);

      // Assert
      expect(spawnerService.spawnAgent).toHaveBeenCalledWith(reviewerRequest);
      expect(result.agentId).toBe(agentId);
    });

    it("should handle tester agent type", async () => {
      // Arrange
      const testerRequest = {
        ...validRequest,
        agentType: "tester" as const,
      };
      const agentId = "agent-tester-123";
      spawnerService.spawnAgent.mockReturnValue({
        agentId,
        state: "spawning",
        spawnedAt: new Date(),
      });
      queueService.addTask.mockResolvedValue(undefined);

      // Act
      const result = await controller.spawn(testerRequest);

      // Assert
      expect(spawnerService.spawnAgent).toHaveBeenCalledWith(testerRequest);
      expect(result.agentId).toBe(agentId);
    });

    it("should handle missing optional skills", async () => {
      // Arrange
      const requestWithoutSkills = {
        taskId: "task-123",
        agentType: "worker" as const,
        context: {
          repository: "https://github.com/org/repo.git",
          branch: "main",
          workItems: ["US-001"],
        },
      };
      const agentId = "agent-abc-123";
      spawnerService.spawnAgent.mockReturnValue({
        agentId,
        state: "spawning",
        spawnedAt: new Date(),
      });
      queueService.addTask.mockResolvedValue(undefined);

      // Act
      const result = await controller.spawn(requestWithoutSkills);

      // Assert
      expect(result.agentId).toBe(agentId);
    });

    it("should throw BadRequestException when taskId is missing", async () => {
      // Arrange
      const invalidRequest = {
        agentType: "worker" as const,
        context: validRequest.context,
      } as unknown as typeof validRequest;

      // Act & Assert
      await expect(controller.spawn(invalidRequest)).rejects.toThrow(BadRequestException);
      expect(spawnerService.spawnAgent).not.toHaveBeenCalled();
      expect(queueService.addTask).not.toHaveBeenCalled();
    });

    it("should throw BadRequestException when agentType is invalid", async () => {
      // Arrange
      const invalidRequest = {
        ...validRequest,
        agentType: "invalid" as unknown as "worker",
      };

      // Act & Assert
      await expect(controller.spawn(invalidRequest)).rejects.toThrow(BadRequestException);
      expect(spawnerService.spawnAgent).not.toHaveBeenCalled();
      expect(queueService.addTask).not.toHaveBeenCalled();
    });

    it("should throw BadRequestException when repository is missing", async () => {
      // Arrange
      const invalidRequest = {
        ...validRequest,
        context: {
          ...validRequest.context,
          repository: "",
        },
      };

      // Act & Assert
      await expect(controller.spawn(invalidRequest)).rejects.toThrow(BadRequestException);
      expect(spawnerService.spawnAgent).not.toHaveBeenCalled();
      expect(queueService.addTask).not.toHaveBeenCalled();
    });

    it("should throw BadRequestException when branch is missing", async () => {
      // Arrange
      const invalidRequest = {
        ...validRequest,
        context: {
          ...validRequest.context,
          branch: "",
        },
      };

      // Act & Assert
      await expect(controller.spawn(invalidRequest)).rejects.toThrow(BadRequestException);
      expect(spawnerService.spawnAgent).not.toHaveBeenCalled();
      expect(queueService.addTask).not.toHaveBeenCalled();
    });

    it("should throw BadRequestException when workItems is empty", async () => {
      // Arrange
      const invalidRequest = {
        ...validRequest,
        context: {
          ...validRequest.context,
          workItems: [],
        },
      };

      // Act & Assert
      await expect(controller.spawn(invalidRequest)).rejects.toThrow(BadRequestException);
      expect(spawnerService.spawnAgent).not.toHaveBeenCalled();
      expect(queueService.addTask).not.toHaveBeenCalled();
    });

    it("should propagate errors from spawner service", async () => {
      // Arrange
      const error = new Error("Spawner failed");
      spawnerService.spawnAgent.mockImplementation(() => {
        throw error;
      });

      // Act & Assert
      await expect(controller.spawn(validRequest)).rejects.toThrow("Spawner failed");
      expect(queueService.addTask).not.toHaveBeenCalled();
    });

    it("should propagate errors from queue service", async () => {
      // Arrange
      const agentId = "agent-abc-123";
      spawnerService.spawnAgent.mockReturnValue({
        agentId,
        state: "spawning",
        spawnedAt: new Date(),
      });
      const error = new Error("Queue failed");
      queueService.addTask.mockRejectedValue(error);

      // Act & Assert
      await expect(controller.spawn(validRequest)).rejects.toThrow("Queue failed");
    });

    it("should use default priority of 5", async () => {
      // Arrange
      const agentId = "agent-abc-123";
      spawnerService.spawnAgent.mockReturnValue({
        agentId,
        state: "spawning",
        spawnedAt: new Date(),
      });
      queueService.addTask.mockResolvedValue(undefined);

      // Act
      await controller.spawn(validRequest);

      // Assert
      expect(queueService.addTask).toHaveBeenCalledWith(validRequest.taskId, validRequest.context, {
        priority: 5,
      });
    });
  });
});