/** * Command Service Tests */ import { describe, it, expect, beforeEach, vi } from "vitest"; import { Test, TestingModule } from "@nestjs/testing"; import { HttpService } from "@nestjs/axios"; import { CommandService } from "./command.service"; import { PrismaService } from "../prisma/prisma.service"; import { FederationService } from "./federation.service"; import { SignatureService } from "./signature.service"; import { FederationConnectionStatus, FederationMessageType, FederationMessageStatus, } from "@prisma/client"; import { of } from "rxjs"; import type { CommandMessage, CommandResponse } from "./types/message.types"; describe("CommandService", () => { let service: CommandService; let prisma: PrismaService; let federationService: FederationService; let signatureService: SignatureService; let httpService: HttpService; const mockWorkspaceId = "workspace-123"; const mockConnectionId = "connection-123"; const mockInstanceId = "instance-456"; const mockRemoteUrl = "https://remote.example.com"; beforeEach(async () => { const module: TestingModule = await Test.createTestingModule({ providers: [ CommandService, { provide: PrismaService, useValue: { federationConnection: { findUnique: vi.fn(), findFirst: vi.fn(), }, federationMessage: { create: vi.fn(), update: vi.fn(), findMany: vi.fn(), findUnique: vi.fn(), findFirst: vi.fn(), }, }, }, { provide: FederationService, useValue: { getInstanceIdentity: vi.fn(), }, }, { provide: SignatureService, useValue: { signMessage: vi.fn(), verifyMessage: vi.fn(), validateTimestamp: vi.fn(), }, }, { provide: HttpService, useValue: { post: vi.fn(), }, }, ], }).compile(); service = module.get(CommandService); prisma = module.get(PrismaService); federationService = module.get(FederationService); signatureService = module.get(SignatureService); httpService = module.get(HttpService); }); describe("sendCommand", () => { it("should send a command to a remote instance", async () => { const commandType = "spawn_agent"; const payload = { agentType: "task_executor" }; const mockConnection = { id: mockConnectionId, workspaceId: mockWorkspaceId, status: FederationConnectionStatus.ACTIVE, remoteUrl: mockRemoteUrl, remoteInstanceId: mockInstanceId, }; const mockIdentity = { instanceId: "local-instance", displayName: "Local Instance", }; const mockMessage = { id: "msg-123", workspaceId: mockWorkspaceId, connectionId: mockConnectionId, messageType: FederationMessageType.COMMAND, messageId: expect.any(String), correlationId: null, query: null, commandType, payload, response: {}, status: FederationMessageStatus.PENDING, error: null, signature: "signature-123", createdAt: new Date(), updatedAt: new Date(), deliveredAt: null, }; vi.spyOn(prisma.federationConnection, "findUnique").mockResolvedValue( mockConnection as never ); vi.spyOn(federationService, "getInstanceIdentity").mockResolvedValue(mockIdentity as never); vi.spyOn(signatureService, "signMessage").mockResolvedValue("signature-123"); vi.spyOn(prisma.federationMessage, "create").mockResolvedValue(mockMessage as never); vi.spyOn(httpService, "post").mockReturnValue(of({} as never)); const result = await service.sendCommand( mockWorkspaceId, mockConnectionId, commandType, payload ); expect(result).toMatchObject({ workspaceId: mockWorkspaceId, connectionId: mockConnectionId, messageType: FederationMessageType.COMMAND, commandType, status: FederationMessageStatus.PENDING, }); expect(httpService.post).toHaveBeenCalledWith( `${mockRemoteUrl}/api/v1/federation/incoming/command`, expect.objectContaining({ messageId: expect.any(String), instanceId: "local-instance", commandType, payload, timestamp: expect.any(Number), signature: "signature-123", }) ); }); it("should throw error if connection not found", async () => { vi.spyOn(prisma.federationConnection, "findUnique").mockResolvedValue(null); await expect( service.sendCommand(mockWorkspaceId, mockConnectionId, "test", {}) ).rejects.toThrow("Connection not found"); }); it("should throw error if connection is not active", async () => { const mockConnection = { id: mockConnectionId, workspaceId: mockWorkspaceId, status: FederationConnectionStatus.SUSPENDED, }; vi.spyOn(prisma.federationConnection, "findUnique").mockResolvedValue( mockConnection as never ); await expect( service.sendCommand(mockWorkspaceId, mockConnectionId, "test", {}) ).rejects.toThrow("Connection is not active"); }); it("should mark command as failed if sending fails", async () => { const mockConnection = { id: mockConnectionId, workspaceId: mockWorkspaceId, status: FederationConnectionStatus.ACTIVE, remoteUrl: mockRemoteUrl, }; const mockIdentity = { instanceId: "local-instance", displayName: "Local Instance", }; const mockMessage = { id: "msg-123", workspaceId: mockWorkspaceId, connectionId: mockConnectionId, messageType: FederationMessageType.COMMAND, messageId: "test-msg-id", correlationId: null, query: null, commandType: "test", payload: {}, response: {}, status: FederationMessageStatus.PENDING, error: null, signature: "signature-123", createdAt: new Date(), updatedAt: new Date(), deliveredAt: null, }; vi.spyOn(prisma.federationConnection, "findUnique").mockResolvedValue( mockConnection as never ); vi.spyOn(federationService, "getInstanceIdentity").mockResolvedValue(mockIdentity as never); vi.spyOn(signatureService, "signMessage").mockResolvedValue("signature-123"); vi.spyOn(prisma.federationMessage, "create").mockResolvedValue(mockMessage as never); vi.spyOn(httpService, "post").mockReturnValue( new (class { subscribe(handlers: { error: (err: Error) => void }) { handlers.error(new Error("Network error")); } })() as never ); vi.spyOn(prisma.federationMessage, "update").mockResolvedValue(mockMessage as never); await expect( service.sendCommand(mockWorkspaceId, mockConnectionId, "test", {}) ).rejects.toThrow("Failed to send command"); expect(prisma.federationMessage.update).toHaveBeenCalledWith({ where: { id: "msg-123" }, data: { status: FederationMessageStatus.FAILED, error: "Network error", }, }); }); }); describe("handleIncomingCommand", () => { it("should process a valid incoming command", async () => { const commandMessage: CommandMessage = { messageId: "cmd-123", instanceId: mockInstanceId, commandType: "spawn_agent", payload: { agentType: "task_executor" }, timestamp: Date.now(), signature: "signature-123", }; const mockConnection = { id: mockConnectionId, remoteInstanceId: mockInstanceId, status: FederationConnectionStatus.ACTIVE, }; const mockIdentity = { instanceId: "local-instance", displayName: "Local Instance", }; vi.spyOn(signatureService, "validateTimestamp").mockReturnValue(true); vi.spyOn(prisma.federationConnection, "findFirst").mockResolvedValue(mockConnection as never); vi.spyOn(signatureService, "verifyMessage").mockResolvedValue({ valid: true, error: null, } as never); vi.spyOn(federationService, "getInstanceIdentity").mockResolvedValue(mockIdentity as never); vi.spyOn(signatureService, "signMessage").mockResolvedValue("response-signature"); const response = await service.handleIncomingCommand(commandMessage); expect(response).toMatchObject({ correlationId: "cmd-123", instanceId: "local-instance", success: true, }); expect(signatureService.validateTimestamp).toHaveBeenCalledWith(commandMessage.timestamp); expect(signatureService.verifyMessage).toHaveBeenCalledWith( expect.objectContaining({ messageId: "cmd-123", instanceId: mockInstanceId, commandType: "spawn_agent", }), "signature-123", mockInstanceId ); }); it("should reject command with invalid timestamp", async () => { const commandMessage: CommandMessage = { messageId: "cmd-123", instanceId: mockInstanceId, commandType: "test", payload: {}, timestamp: Date.now() - 1000000, signature: "signature-123", }; vi.spyOn(signatureService, "validateTimestamp").mockReturnValue(false); await expect(service.handleIncomingCommand(commandMessage)).rejects.toThrow( "Command timestamp is outside acceptable range" ); }); it("should reject command if no connection found", async () => { const commandMessage: CommandMessage = { messageId: "cmd-123", instanceId: mockInstanceId, commandType: "test", payload: {}, timestamp: Date.now(), signature: "signature-123", }; vi.spyOn(signatureService, "validateTimestamp").mockReturnValue(true); vi.spyOn(prisma.federationConnection, "findFirst").mockResolvedValue(null); await expect(service.handleIncomingCommand(commandMessage)).rejects.toThrow( "No connection found for remote instance" ); }); it("should reject command with invalid signature", async () => { const commandMessage: CommandMessage = { messageId: "cmd-123", instanceId: mockInstanceId, commandType: "test", payload: {}, timestamp: Date.now(), signature: "invalid-signature", }; const mockConnection = { id: mockConnectionId, remoteInstanceId: mockInstanceId, status: FederationConnectionStatus.ACTIVE, }; vi.spyOn(signatureService, "validateTimestamp").mockReturnValue(true); vi.spyOn(prisma.federationConnection, "findFirst").mockResolvedValue(mockConnection as never); vi.spyOn(signatureService, "verifyMessage").mockResolvedValue({ valid: false, error: "Invalid signature", } as never); await expect(service.handleIncomingCommand(commandMessage)).rejects.toThrow( "Invalid signature" ); }); }); describe("processCommandResponse", () => { it("should process a successful command response", async () => { const response: CommandResponse = { messageId: "resp-123", correlationId: "cmd-123", instanceId: mockInstanceId, success: true, data: { result: "success" }, timestamp: Date.now(), signature: "signature-123", }; const mockMessage = { id: "msg-123", workspaceId: mockWorkspaceId, connectionId: mockConnectionId, messageType: FederationMessageType.COMMAND, messageId: "cmd-123", correlationId: null, query: null, commandType: "test", payload: {}, response: {}, status: FederationMessageStatus.PENDING, error: null, signature: "signature-123", createdAt: new Date(), updatedAt: new Date(), deliveredAt: null, }; vi.spyOn(signatureService, "validateTimestamp").mockReturnValue(true); vi.spyOn(prisma.federationMessage, "findFirst").mockResolvedValue(mockMessage as never); vi.spyOn(signatureService, "verifyMessage").mockResolvedValue({ valid: true, error: null, } as never); vi.spyOn(prisma.federationMessage, "update").mockResolvedValue(mockMessage as never); await service.processCommandResponse(response); expect(prisma.federationMessage.update).toHaveBeenCalledWith({ where: { id: "msg-123" }, data: { status: FederationMessageStatus.DELIVERED, deliveredAt: expect.any(Date), response: { result: "success" }, }, }); }); it("should handle failed command response", async () => { const response: CommandResponse = { messageId: "resp-123", correlationId: "cmd-123", instanceId: mockInstanceId, success: false, error: "Command execution failed", timestamp: Date.now(), signature: "signature-123", }; const mockMessage = { id: "msg-123", messageType: FederationMessageType.COMMAND, messageId: "cmd-123", }; vi.spyOn(signatureService, "validateTimestamp").mockReturnValue(true); vi.spyOn(prisma.federationMessage, "findFirst").mockResolvedValue(mockMessage as never); vi.spyOn(signatureService, "verifyMessage").mockResolvedValue({ valid: true, error: null, } as never); vi.spyOn(prisma.federationMessage, "update").mockResolvedValue(mockMessage as never); await service.processCommandResponse(response); expect(prisma.federationMessage.update).toHaveBeenCalledWith({ where: { id: "msg-123" }, data: { status: FederationMessageStatus.FAILED, deliveredAt: expect.any(Date), error: "Command execution failed", }, }); }); it("should reject response with invalid timestamp", async () => { const response: CommandResponse = { messageId: "resp-123", correlationId: "cmd-123", instanceId: mockInstanceId, success: true, timestamp: Date.now() - 1000000, signature: "signature-123", }; vi.spyOn(signatureService, "validateTimestamp").mockReturnValue(false); await expect(service.processCommandResponse(response)).rejects.toThrow( "Response timestamp is outside acceptable range" ); }); }); describe("getCommandMessages", () => { it("should return all command messages for a workspace", async () => { const mockMessages = [ { id: "msg-1", workspaceId: mockWorkspaceId, connectionId: mockConnectionId, messageType: FederationMessageType.COMMAND, messageId: "cmd-1", correlationId: null, query: null, commandType: "test", payload: {}, response: {}, status: FederationMessageStatus.DELIVERED, error: null, signature: "sig-1", createdAt: new Date(), updatedAt: new Date(), deliveredAt: new Date(), }, ]; vi.spyOn(prisma.federationMessage, "findMany").mockResolvedValue(mockMessages as never); const result = await service.getCommandMessages(mockWorkspaceId); expect(result).toHaveLength(1); expect(result[0]).toMatchObject({ workspaceId: mockWorkspaceId, messageType: FederationMessageType.COMMAND, commandType: "test", }); }); it("should filter command messages by status", async () => { const mockMessages = [ { id: "msg-1", workspaceId: mockWorkspaceId, connectionId: mockConnectionId, messageType: FederationMessageType.COMMAND, messageId: "cmd-1", correlationId: null, query: null, commandType: "test", payload: {}, response: {}, status: FederationMessageStatus.PENDING, error: null, signature: "sig-1", createdAt: new Date(), updatedAt: new Date(), deliveredAt: null, }, ]; vi.spyOn(prisma.federationMessage, "findMany").mockResolvedValue(mockMessages as never); await service.getCommandMessages(mockWorkspaceId, FederationMessageStatus.PENDING); expect(prisma.federationMessage.findMany).toHaveBeenCalledWith({ where: { workspaceId: mockWorkspaceId, messageType: FederationMessageType.COMMAND, status: FederationMessageStatus.PENDING, }, orderBy: { createdAt: "desc" }, }); }); }); describe("getCommandMessage", () => { it("should return a single command message", async () => { const mockMessage = { id: "msg-1", workspaceId: mockWorkspaceId, connectionId: mockConnectionId, messageType: FederationMessageType.COMMAND, messageId: "cmd-1", correlationId: null, query: null, commandType: "test", payload: { key: "value" }, response: {}, status: FederationMessageStatus.DELIVERED, error: null, signature: "sig-1", createdAt: new Date(), updatedAt: new Date(), deliveredAt: new Date(), }; vi.spyOn(prisma.federationMessage, "findUnique").mockResolvedValue(mockMessage as never); const result = await service.getCommandMessage(mockWorkspaceId, "msg-1"); expect(result).toMatchObject({ id: "msg-1", workspaceId: mockWorkspaceId, commandType: "test", payload: { key: "value" }, }); }); it("should throw error if command message not found", async () => { vi.spyOn(prisma.federationMessage, "findUnique").mockResolvedValue(null); await expect(service.getCommandMessage(mockWorkspaceId, "invalid-id")).rejects.toThrow( "Command message not found" ); }); }); });