stack/apps/api/src/federation/command.service.spec.ts

/**
 * 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>(CommandService);
    prisma = module.get<PrismaService>(PrismaService);
    federationService = module.get<FederationService>(FederationService);
    signatureService = module.get<SignatureService>(SignatureService);
    httpService = module.get<HttpService>(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"
      );
    });
  });
});