/** * Tests for Federation Agent Service */ import { describe, it, expect, beforeEach, vi } from "vitest"; import { Test, TestingModule } from "@nestjs/testing"; import { HttpService } from "@nestjs/axios"; import { ConfigService } from "@nestjs/config"; import { FederationAgentService } from "./federation-agent.service"; import { CommandService } from "./command.service"; import { PrismaService } from "../prisma/prisma.service"; import { FederationConnectionStatus } from "@prisma/client"; import { of, throwError } from "rxjs"; import type { SpawnAgentCommandPayload, AgentStatusCommandPayload, KillAgentCommandPayload, SpawnAgentResponseData, AgentStatusResponseData, KillAgentResponseData, } from "./types/federation-agent.types"; describe("FederationAgentService", () => { let service: FederationAgentService; let commandService: ReturnType>; let prisma: ReturnType>; let httpService: ReturnType>; let configService: ReturnType>; const mockWorkspaceId = "workspace-1"; const mockConnectionId = "connection-1"; const mockAgentId = "agent-123"; const mockTaskId = "task-456"; const mockOrchestratorUrl = "http://localhost:3001"; beforeEach(async () => { const mockCommandService = { sendCommand: vi.fn(), }; const mockPrisma = { federationConnection: { findUnique: vi.fn(), findFirst: vi.fn(), }, }; const mockHttpService = { post: vi.fn(), get: vi.fn(), }; const mockConfigService = { get: vi.fn((key: string) => { if (key === "orchestrator.url") { return mockOrchestratorUrl; } return undefined; }), }; const module: TestingModule = await Test.createTestingModule({ providers: [ FederationAgentService, { provide: CommandService, useValue: mockCommandService }, { provide: PrismaService, useValue: mockPrisma }, { provide: HttpService, useValue: mockHttpService }, { provide: ConfigService, useValue: mockConfigService }, ], }).compile(); service = module.get(FederationAgentService); commandService = module.get(CommandService); prisma = module.get(PrismaService); httpService = module.get(HttpService); configService = module.get(ConfigService); }); it("should be defined", () => { expect(service).toBeDefined(); }); describe("spawnAgentOnRemote", () => { const spawnPayload: SpawnAgentCommandPayload = { taskId: mockTaskId, agentType: "worker", context: { repository: "git.example.com/org/repo", branch: "main", workItems: ["item-1"], }, }; const mockConnection = { id: mockConnectionId, workspaceId: mockWorkspaceId, remoteInstanceId: "remote-instance-1", remoteUrl: "https://remote.example.com", status: FederationConnectionStatus.ACTIVE, }; it("should spawn agent on remote instance", async () => { prisma.federationConnection.findUnique.mockResolvedValue(mockConnection as never); const mockCommandResponse = { id: "msg-1", workspaceId: mockWorkspaceId, connectionId: mockConnectionId, messageType: "COMMAND" as never, messageId: "msg-uuid", commandType: "agent.spawn", payload: spawnPayload as never, response: { agentId: mockAgentId, status: "spawning", spawnedAt: "2026-02-03T14:30:00Z", } as never, status: "DELIVERED" as never, createdAt: new Date(), updatedAt: new Date(), }; commandService.sendCommand.mockResolvedValue(mockCommandResponse as never); const result = await service.spawnAgentOnRemote( mockWorkspaceId, mockConnectionId, spawnPayload ); expect(prisma.federationConnection.findUnique).toHaveBeenCalledWith({ where: { id: mockConnectionId, workspaceId: mockWorkspaceId }, }); expect(commandService.sendCommand).toHaveBeenCalledWith( mockWorkspaceId, mockConnectionId, "agent.spawn", spawnPayload ); expect(result).toEqual(mockCommandResponse); }); it("should throw error if connection not found", async () => { prisma.federationConnection.findUnique.mockResolvedValue(null); await expect( service.spawnAgentOnRemote(mockWorkspaceId, mockConnectionId, spawnPayload) ).rejects.toThrow("Connection not found"); expect(commandService.sendCommand).not.toHaveBeenCalled(); }); it("should throw error if connection not active", async () => { const inactiveConnection = { ...mockConnection, status: FederationConnectionStatus.DISCONNECTED, }; prisma.federationConnection.findUnique.mockResolvedValue(inactiveConnection as never); await expect( service.spawnAgentOnRemote(mockWorkspaceId, mockConnectionId, spawnPayload) ).rejects.toThrow("Connection is not active"); expect(commandService.sendCommand).not.toHaveBeenCalled(); }); }); describe("getAgentStatus", () => { const statusPayload: AgentStatusCommandPayload = { agentId: mockAgentId, }; const mockConnection = { id: mockConnectionId, workspaceId: mockWorkspaceId, remoteInstanceId: "remote-instance-1", remoteUrl: "https://remote.example.com", status: FederationConnectionStatus.ACTIVE, }; it("should get agent status from remote instance", async () => { prisma.federationConnection.findUnique.mockResolvedValue(mockConnection as never); const mockCommandResponse = { id: "msg-2", workspaceId: mockWorkspaceId, connectionId: mockConnectionId, messageType: "COMMAND" as never, messageId: "msg-uuid-2", commandType: "agent.status", payload: statusPayload as never, response: { agentId: mockAgentId, taskId: mockTaskId, status: "running", spawnedAt: "2026-02-03T14:30:00Z", startedAt: "2026-02-03T14:30:05Z", } as never, status: "DELIVERED" as never, createdAt: new Date(), updatedAt: new Date(), }; commandService.sendCommand.mockResolvedValue(mockCommandResponse as never); const result = await service.getAgentStatus(mockWorkspaceId, mockConnectionId, mockAgentId); expect(commandService.sendCommand).toHaveBeenCalledWith( mockWorkspaceId, mockConnectionId, "agent.status", statusPayload ); expect(result).toEqual(mockCommandResponse); }); }); describe("killAgentOnRemote", () => { const killPayload: KillAgentCommandPayload = { agentId: mockAgentId, }; const mockConnection = { id: mockConnectionId, workspaceId: mockWorkspaceId, remoteInstanceId: "remote-instance-1", remoteUrl: "https://remote.example.com", status: FederationConnectionStatus.ACTIVE, }; it("should kill agent on remote instance", async () => { prisma.federationConnection.findUnique.mockResolvedValue(mockConnection as never); const mockCommandResponse = { id: "msg-3", workspaceId: mockWorkspaceId, connectionId: mockConnectionId, messageType: "COMMAND" as never, messageId: "msg-uuid-3", commandType: "agent.kill", payload: killPayload as never, response: { agentId: mockAgentId, status: "killed", killedAt: "2026-02-03T14:35:00Z", } as never, status: "DELIVERED" as never, createdAt: new Date(), updatedAt: new Date(), }; commandService.sendCommand.mockResolvedValue(mockCommandResponse as never); const result = await service.killAgentOnRemote( mockWorkspaceId, mockConnectionId, mockAgentId ); expect(commandService.sendCommand).toHaveBeenCalledWith( mockWorkspaceId, mockConnectionId, "agent.kill", killPayload ); expect(result).toEqual(mockCommandResponse); }); }); describe("handleAgentCommand", () => { const mockConnection = { id: mockConnectionId, workspaceId: mockWorkspaceId, remoteInstanceId: "remote-instance-1", remoteUrl: "https://remote.example.com", status: FederationConnectionStatus.ACTIVE, }; it("should handle agent.spawn command", async () => { const spawnPayload: SpawnAgentCommandPayload = { taskId: mockTaskId, agentType: "worker", context: { repository: "git.example.com/org/repo", branch: "main", workItems: ["item-1"], }, }; prisma.federationConnection.findFirst.mockResolvedValue(mockConnection as never); const mockOrchestratorResponse = { agentId: mockAgentId, status: "spawning", }; httpService.post.mockReturnValue( of({ data: mockOrchestratorResponse, status: 200, statusText: "OK", headers: {}, config: {} as never, }) as never ); const result = await service.handleAgentCommand( "remote-instance-1", "agent.spawn", spawnPayload ); expect(httpService.post).toHaveBeenCalledWith( `${mockOrchestratorUrl}/agents/spawn`, expect.objectContaining({ taskId: mockTaskId, agentType: "worker", }) ); expect(result.success).toBe(true); expect(result.data).toEqual({ agentId: mockAgentId, status: "spawning", spawnedAt: expect.any(String), }); }); it("should handle agent.status command", async () => { const statusPayload: AgentStatusCommandPayload = { agentId: mockAgentId, }; prisma.federationConnection.findFirst.mockResolvedValue(mockConnection as never); const mockOrchestratorResponse = { agentId: mockAgentId, taskId: mockTaskId, status: "running", spawnedAt: "2026-02-03T14:30:00Z", startedAt: "2026-02-03T14:30:05Z", }; httpService.get.mockReturnValue( of({ data: mockOrchestratorResponse, status: 200, statusText: "OK", headers: {}, config: {} as never, }) as never ); const result = await service.handleAgentCommand( "remote-instance-1", "agent.status", statusPayload ); expect(httpService.get).toHaveBeenCalledWith( `${mockOrchestratorUrl}/agents/${mockAgentId}/status` ); expect(result.success).toBe(true); expect(result.data).toEqual(mockOrchestratorResponse); }); it("should handle agent.kill command", async () => { const killPayload: KillAgentCommandPayload = { agentId: mockAgentId, }; prisma.federationConnection.findFirst.mockResolvedValue(mockConnection as never); const mockOrchestratorResponse = { message: `Agent ${mockAgentId} killed successfully`, }; httpService.post.mockReturnValue( of({ data: mockOrchestratorResponse, status: 200, statusText: "OK", headers: {}, config: {} as never, }) as never ); const result = await service.handleAgentCommand( "remote-instance-1", "agent.kill", killPayload ); expect(httpService.post).toHaveBeenCalledWith( `${mockOrchestratorUrl}/agents/${mockAgentId}/kill`, {} ); expect(result.success).toBe(true); expect(result.data).toEqual({ agentId: mockAgentId, status: "killed", killedAt: expect.any(String), }); }); it("should return error for unknown command type", async () => { prisma.federationConnection.findFirst.mockResolvedValue(mockConnection as never); const result = await service.handleAgentCommand("remote-instance-1", "agent.unknown", {}); expect(result.success).toBe(false); expect(result.error).toContain("Unknown agent command type: agent.unknown"); }); it("should throw error if connection not found", async () => { prisma.federationConnection.findFirst.mockResolvedValue(null); await expect( service.handleAgentCommand("remote-instance-1", "agent.spawn", {}) ).rejects.toThrow("No connection found for remote instance"); }); it("should handle orchestrator errors", async () => { const spawnPayload: SpawnAgentCommandPayload = { taskId: mockTaskId, agentType: "worker", context: { repository: "git.example.com/org/repo", branch: "main", workItems: ["item-1"], }, }; prisma.federationConnection.findFirst.mockResolvedValue(mockConnection as never); httpService.post.mockReturnValue( throwError(() => new Error("Orchestrator connection failed")) as never ); const result = await service.handleAgentCommand( "remote-instance-1", "agent.spawn", spawnPayload ); expect(result.success).toBe(false); expect(result.error).toContain("Orchestrator connection failed"); }); }); });