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

/**
 * Connection Service Tests
 *
 * Tests for federation connection management.
 */

import { describe, it, expect, beforeEach, vi } from "vitest";
import { Test, TestingModule } from "@nestjs/testing";
import { HttpService } from "@nestjs/axios";
import { ConnectionService } from "./connection.service";
import { FederationService } from "./federation.service";
import { SignatureService } from "./signature.service";
import { FederationAuditService } from "./audit.service";
import { PrismaService } from "../prisma/prisma.service";
import { FederationConnectionStatus } from "@prisma/client";
import { FederationConnection } from "@prisma/client";
import { of, throwError } from "rxjs";
import type { AxiosResponse } from "axios";

describe("ConnectionService", () => {
  let service: ConnectionService;
  let prismaService: PrismaService;
  let federationService: FederationService;
  let signatureService: SignatureService;
  let httpService: HttpService;
  let auditService: FederationAuditService;

  const mockWorkspaceId = "workspace-123";
  const mockRemoteUrl = "https://remote.example.com";
  const mockInstanceIdentity = {
    id: "local-id",
    instanceId: "local-instance-123",
    name: "Local Instance",
    url: "https://local.example.com",
    publicKey: "-----BEGIN PUBLIC KEY-----\nLOCAL\n-----END PUBLIC KEY-----",
    privateKey: "-----BEGIN PRIVATE KEY-----\nLOCAL\n-----END PRIVATE KEY-----",
    capabilities: {
      supportsQuery: true,
      supportsCommand: true,
      protocolVersion: "1.0",
    },
    metadata: {},
    createdAt: new Date(),
    updatedAt: new Date(),
  };

  const mockRemoteIdentity = {
    id: "remote-id",
    instanceId: "remote-instance-456",
    name: "Remote Instance",
    url: mockRemoteUrl,
    publicKey: "-----BEGIN PUBLIC KEY-----\nREMOTE\n-----END PUBLIC KEY-----",
    capabilities: {
      supportsQuery: true,
      protocolVersion: "1.0",
    },
    metadata: {},
    createdAt: new Date(),
    updatedAt: new Date(),
  };

  const mockConnection: FederationConnection = {
    id: "conn-123",
    workspaceId: mockWorkspaceId,
    remoteInstanceId: mockRemoteIdentity.instanceId,
    remoteUrl: mockRemoteUrl,
    remotePublicKey: mockRemoteIdentity.publicKey,
    remoteCapabilities: mockRemoteIdentity.capabilities,
    status: FederationConnectionStatus.PENDING,
    metadata: {},
    createdAt: new Date(),
    updatedAt: new Date(),
    connectedAt: null,
    disconnectedAt: null,
  };

  beforeEach(async () => {
    const module: TestingModule = await Test.createTestingModule({
      providers: [
        ConnectionService,
        {
          provide: PrismaService,
          useValue: {
            federationConnection: {
              create: vi.fn(),
              findFirst: vi.fn(),
              findUnique: vi.fn(),
              findMany: vi.fn(),
              update: vi.fn(),
              delete: vi.fn(),
              count: vi.fn(),
            },
          },
        },
        {
          provide: FederationService,
          useValue: {
            getInstanceIdentity: vi.fn().mockResolvedValue(mockInstanceIdentity),
            getPublicIdentity: vi.fn().mockResolvedValue(mockInstanceIdentity),
          },
        },
        {
          provide: SignatureService,
          useValue: {
            signMessage: vi.fn().mockResolvedValue("mock-signature"),
            verifyConnectionRequest: vi.fn().mockResolvedValue({ valid: true }),
          },
        },
        {
          provide: HttpService,
          useValue: {
            get: vi.fn(),
            post: vi.fn(),
          },
        },
        {
          provide: FederationAuditService,
          useValue: {
            logIncomingConnectionAttempt: vi.fn(),
            logIncomingConnectionCreated: vi.fn(),
            logIncomingConnectionRejected: vi.fn(),
          },
        },
      ],
    }).compile();

    service = module.get<ConnectionService>(ConnectionService);
    prismaService = module.get<PrismaService>(PrismaService);
    federationService = module.get<FederationService>(FederationService);
    signatureService = module.get<SignatureService>(SignatureService);
    httpService = module.get<HttpService>(HttpService);
    auditService = module.get<FederationAuditService>(FederationAuditService);
  });

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

  describe("initiateConnection", () => {
    it("should throw error if workspace has reached connection limit", async () => {
      const existingConnections = Array.from({ length: 100 }, (_, i) => ({
        ...mockConnection,
        id: `conn-${i}`,
      }));

      vi.spyOn(prismaService.federationConnection, "count").mockResolvedValue(100);

      await expect(service.initiateConnection(mockWorkspaceId, mockRemoteUrl)).rejects.toThrow(
        "Connection limit reached for workspace. Maximum 100 connections allowed per workspace."
      );
    });

    it("should reject connection to instance with incompatible protocol version", async () => {
      const incompatibleRemoteIdentity = {
        ...mockRemoteIdentity,
        capabilities: {
          ...mockRemoteIdentity.capabilities,
          protocolVersion: "2.0",
        },
      };

      const mockAxiosResponse: AxiosResponse = {
        data: incompatibleRemoteIdentity,
        status: 200,
        statusText: "OK",
        headers: {},
        config: {} as never,
      };

      vi.spyOn(prismaService.federationConnection, "count").mockResolvedValue(5);
      vi.spyOn(httpService, "get").mockReturnValue(of(mockAxiosResponse));

      await expect(service.initiateConnection(mockWorkspaceId, mockRemoteUrl)).rejects.toThrow(
        "Incompatible protocol version. Expected 1.0, received 2.0"
      );
    });

    it("should create a pending connection", async () => {
      const mockAxiosResponse: AxiosResponse = {
        data: mockRemoteIdentity,
        status: 200,
        statusText: "OK",
        headers: {},
        config: {} as never,
      };

      vi.spyOn(prismaService.federationConnection, "count").mockResolvedValue(5);
      vi.spyOn(httpService, "get").mockReturnValue(of(mockAxiosResponse));
      vi.spyOn(httpService, "post").mockReturnValue(
        of({ data: { accepted: true } } as AxiosResponse)
      );
      vi.spyOn(prismaService.federationConnection, "create").mockResolvedValue(mockConnection);

      const result = await service.initiateConnection(mockWorkspaceId, mockRemoteUrl);

      expect(result).toBeDefined();
      expect(result.status).toBe(FederationConnectionStatus.PENDING);
      expect(result.remoteUrl).toBe(mockRemoteUrl);
      expect(prismaService.federationConnection.create).toHaveBeenCalledWith(
        expect.objectContaining({
          data: expect.objectContaining({
            workspaceId: mockWorkspaceId,
            remoteUrl: mockRemoteUrl,
            status: FederationConnectionStatus.PENDING,
          }),
        })
      );
    });

    it("should fetch remote instance identity", async () => {
      const mockAxiosResponse: AxiosResponse = {
        data: mockRemoteIdentity,
        status: 200,
        statusText: "OK",
        headers: {},
        config: {} as never,
      };

      vi.spyOn(prismaService.federationConnection, "count").mockResolvedValue(5);
      vi.spyOn(httpService, "get").mockReturnValue(of(mockAxiosResponse));
      vi.spyOn(httpService, "post").mockReturnValue(
        of({ data: { accepted: true } } as AxiosResponse)
      );
      vi.spyOn(prismaService.federationConnection, "create").mockResolvedValue(mockConnection);

      await service.initiateConnection(mockWorkspaceId, mockRemoteUrl);

      expect(httpService.get).toHaveBeenCalledWith(`${mockRemoteUrl}/api/v1/federation/instance`);
    });

    it("should throw error if remote instance not reachable", async () => {
      vi.spyOn(httpService, "get").mockReturnValue(throwError(() => new Error("Network error")));

      await expect(service.initiateConnection(mockWorkspaceId, mockRemoteUrl)).rejects.toThrow();
    });

    it("should send signed connection request", async () => {
      const mockAxiosResponse: AxiosResponse = {
        data: mockRemoteIdentity,
        status: 200,
        statusText: "OK",
        headers: {},
        config: {} as never,
      };

      vi.spyOn(prismaService.federationConnection, "count").mockResolvedValue(5);
      const postSpy = vi
        .spyOn(httpService, "post")
        .mockReturnValue(of({ data: { accepted: true } } as AxiosResponse));
      vi.spyOn(httpService, "get").mockReturnValue(of(mockAxiosResponse));
      vi.spyOn(prismaService.federationConnection, "create").mockResolvedValue(mockConnection);

      await service.initiateConnection(mockWorkspaceId, mockRemoteUrl);

      expect(postSpy).toHaveBeenCalledWith(
        `${mockRemoteUrl}/api/v1/federation/incoming/connect`,
        expect.objectContaining({
          instanceId: mockInstanceIdentity.instanceId,
          instanceUrl: mockInstanceIdentity.url,
          publicKey: mockInstanceIdentity.publicKey,
          signature: "mock-signature",
        })
      );
    });

    it("should delete connection and throw error if request fails", async () => {
      const mockAxiosResponse: AxiosResponse = {
        data: mockRemoteIdentity,
        status: 200,
        statusText: "OK",
        headers: {},
        config: {} as never,
      };

      vi.spyOn(prismaService.federationConnection, "count").mockResolvedValue(5);
      vi.spyOn(httpService, "get").mockReturnValue(of(mockAxiosResponse));
      vi.spyOn(httpService, "post").mockReturnValue(
        throwError(() => new Error("Connection refused"))
      );
      const createSpy = vi
        .spyOn(prismaService.federationConnection, "create")
        .mockResolvedValue(mockConnection);
      const deleteSpy = vi
        .spyOn(prismaService.federationConnection, "delete")
        .mockResolvedValue(mockConnection);

      await expect(service.initiateConnection(mockWorkspaceId, mockRemoteUrl)).rejects.toThrow(
        "Failed to initiate connection"
      );

      expect(createSpy).toHaveBeenCalled();
      expect(deleteSpy).toHaveBeenCalledWith({
        where: { id: mockConnection.id },
      });
    });
  });

  describe("acceptConnection", () => {
    it("should update connection status to ACTIVE", async () => {
      vi.spyOn(prismaService.federationConnection, "findFirst").mockResolvedValue(mockConnection);
      vi.spyOn(prismaService.federationConnection, "update").mockReturnValue({
        ...mockConnection,
        status: FederationConnectionStatus.ACTIVE,
        connectedAt: new Date(),
      });

      const result = await service.acceptConnection(mockWorkspaceId, mockConnection.id);

      expect(result.status).toBe(FederationConnectionStatus.ACTIVE);
      expect(result.connectedAt).toBeDefined();
      expect(prismaService.federationConnection.update).toHaveBeenCalledWith(
        expect.objectContaining({
          where: expect.objectContaining({
            id: mockConnection.id,
          }),
          data: expect.objectContaining({
            status: FederationConnectionStatus.ACTIVE,
            connectedAt: expect.any(Date),
          }),
        })
      );
    });

    it("should throw error if connection not found", async () => {
      vi.spyOn(prismaService.federationConnection, "findFirst").mockResolvedValue(null);

      await expect(service.acceptConnection(mockWorkspaceId, "non-existent-id")).rejects.toThrow(
        "Connection not found"
      );
    });

    it("should enforce workspace isolation", async () => {
      vi.spyOn(prismaService.federationConnection, "findFirst").mockResolvedValue(null);

      await expect(
        service.acceptConnection("different-workspace", mockConnection.id)
      ).rejects.toThrow("Connection not found");
    });
  });

  describe("rejectConnection", () => {
    it("should update connection status to DISCONNECTED", async () => {
      vi.spyOn(prismaService.federationConnection, "findFirst").mockResolvedValue(mockConnection);
      vi.spyOn(prismaService.federationConnection, "update").mockReturnValue({
        ...mockConnection,
        status: FederationConnectionStatus.DISCONNECTED,
        metadata: { rejectionReason: "Not approved" },
      });

      const result = await service.rejectConnection(
        mockWorkspaceId,
        mockConnection.id,
        "Not approved"
      );

      expect(result.status).toBe(FederationConnectionStatus.DISCONNECTED);
      expect(result.metadata).toHaveProperty("rejectionReason", "Not approved");
    });

    it("should throw error if connection not found", async () => {
      vi.spyOn(prismaService.federationConnection, "findFirst").mockResolvedValue(null);

      await expect(
        service.rejectConnection(mockWorkspaceId, "non-existent-id", "Reason")
      ).rejects.toThrow("Connection not found");
    });
  });

  describe("disconnect", () => {
    const activeConnection: FederationConnection = {
      ...mockConnection,
      status: FederationConnectionStatus.ACTIVE,
      connectedAt: new Date(),
    };

    it("should disconnect active connection", async () => {
      vi.spyOn(prismaService.federationConnection, "findFirst").mockResolvedValue(activeConnection);
      vi.spyOn(prismaService.federationConnection, "update").mockReturnValue({
        ...activeConnection,
        status: FederationConnectionStatus.DISCONNECTED,
        disconnectedAt: new Date(),
      });

      const result = await service.disconnect(
        mockWorkspaceId,
        mockConnection.id,
        "Manual disconnect"
      );

      expect(result.status).toBe(FederationConnectionStatus.DISCONNECTED);
      expect(result.disconnectedAt).toBeDefined();
    });

    it("should store disconnection reason in metadata", async () => {
      vi.spyOn(prismaService.federationConnection, "findFirst").mockResolvedValue(activeConnection);
      vi.spyOn(prismaService.federationConnection, "update").mockReturnValue({
        ...activeConnection,
        status: FederationConnectionStatus.DISCONNECTED,
        disconnectedAt: new Date(),
        metadata: { disconnectReason: "Test reason" },
      });

      const result = await service.disconnect(mockWorkspaceId, mockConnection.id, "Test reason");

      expect(result.metadata).toHaveProperty("disconnectReason", "Test reason");
    });
  });

  describe("getConnections", () => {
    it("should list all connections for workspace", async () => {
      const connections = [mockConnection];
      vi.spyOn(prismaService.federationConnection, "findMany").mockResolvedValue(connections);

      const result = await service.getConnections(mockWorkspaceId);

      expect(result).toEqual(connections);
      expect(prismaService.federationConnection.findMany).toHaveBeenCalledWith(
        expect.objectContaining({
          where: { workspaceId: mockWorkspaceId },
        })
      );
    });

    it("should filter by status if provided", async () => {
      const connections = [mockConnection];
      vi.spyOn(prismaService.federationConnection, "findMany").mockResolvedValue(connections);

      await service.getConnections(mockWorkspaceId, FederationConnectionStatus.ACTIVE);

      expect(prismaService.federationConnection.findMany).toHaveBeenCalledWith(
        expect.objectContaining({
          where: {
            workspaceId: mockWorkspaceId,
            status: FederationConnectionStatus.ACTIVE,
          },
        })
      );
    });
  });

  describe("getConnection", () => {
    it("should return connection details", async () => {
      vi.spyOn(prismaService.federationConnection, "findFirst").mockResolvedValue(mockConnection);

      const result = await service.getConnection(mockWorkspaceId, mockConnection.id);

      expect(result).toEqual(mockConnection);
    });

    it("should throw error if connection not found", async () => {
      vi.spyOn(prismaService.federationConnection, "findFirst").mockResolvedValue(null);

      await expect(service.getConnection(mockWorkspaceId, "non-existent-id")).rejects.toThrow(
        "Connection not found"
      );
    });

    it("should enforce workspace isolation", async () => {
      vi.spyOn(prismaService.federationConnection, "findFirst").mockResolvedValue(null);

      await expect(service.getConnection("different-workspace", mockConnection.id)).rejects.toThrow(
        "Connection not found"
      );
    });
  });

  describe("handleIncomingConnectionRequest", () => {
    const mockRequest = {
      instanceId: mockRemoteIdentity.instanceId,
      instanceUrl: mockRemoteIdentity.url,
      publicKey: mockRemoteIdentity.publicKey,
      capabilities: mockRemoteIdentity.capabilities,
      timestamp: Date.now(),
      signature: "valid-signature",
    };

    it("should reject request with incompatible protocol version", async () => {
      const incompatibleRequest = {
        ...mockRequest,
        capabilities: {
          ...mockRemoteIdentity.capabilities,
          protocolVersion: "2.0",
        },
      };

      vi.spyOn(signatureService, "verifyConnectionRequest").mockResolvedValue({ valid: true });

      await expect(
        service.handleIncomingConnectionRequest(mockWorkspaceId, incompatibleRequest)
      ).rejects.toThrow("Incompatible protocol version. Expected 1.0, received 2.0");
    });

    it("should accept request with compatible protocol version", async () => {
      const compatibleRequest = {
        ...mockRequest,
        capabilities: {
          ...mockRemoteIdentity.capabilities,
          protocolVersion: "1.0",
        },
      };

      vi.spyOn(signatureService, "verifyConnectionRequest").mockResolvedValue({ valid: true });
      vi.spyOn(prismaService.federationConnection, "create").mockResolvedValue(mockConnection);

      const result = await service.handleIncomingConnectionRequest(
        mockWorkspaceId,
        compatibleRequest
      );

      expect(result.status).toBe(FederationConnectionStatus.PENDING);
    });

    it("should validate connection request signature", async () => {
      const verifySpy = vi.spyOn(signatureService, "verifyConnectionRequest");
      vi.spyOn(prismaService.federationConnection, "create").mockResolvedValue(mockConnection);

      await service.handleIncomingConnectionRequest(mockWorkspaceId, mockRequest);

      expect(verifySpy).toHaveBeenCalledWith(mockRequest);
    });

    it("should create pending connection for valid request", async () => {
      vi.spyOn(signatureService, "verifyConnectionRequest").mockResolvedValue({ valid: true });
      vi.spyOn(prismaService.federationConnection, "create").mockResolvedValue(mockConnection);

      const result = await service.handleIncomingConnectionRequest(mockWorkspaceId, mockRequest);

      expect(result.status).toBe(FederationConnectionStatus.PENDING);
      expect(prismaService.federationConnection.create).toHaveBeenCalled();
    });

    it("should reject request with invalid signature", async () => {
      vi.spyOn(signatureService, "verifyConnectionRequest").mockResolvedValue({
        valid: false,
        error: "Invalid signature",
      });

      await expect(
        service.handleIncomingConnectionRequest(mockWorkspaceId, mockRequest)
      ).rejects.toThrow("Invalid connection request signature");
    });

    it("should log incoming connection attempt", async () => {
      vi.spyOn(signatureService, "verifyConnectionRequest").mockResolvedValue({ valid: true });
      vi.spyOn(prismaService.federationConnection, "create").mockResolvedValue(mockConnection);
      const auditSpy = vi.spyOn(auditService, "logIncomingConnectionAttempt");

      await service.handleIncomingConnectionRequest(mockWorkspaceId, mockRequest);

      expect(auditSpy).toHaveBeenCalledWith({
        workspaceId: mockWorkspaceId,
        remoteInstanceId: mockRequest.instanceId,
        remoteUrl: mockRequest.instanceUrl,
        timestamp: mockRequest.timestamp,
      });
    });

    it("should log connection created on success", async () => {
      vi.spyOn(signatureService, "verifyConnectionRequest").mockResolvedValue({ valid: true });
      vi.spyOn(prismaService.federationConnection, "create").mockResolvedValue(mockConnection);
      const auditSpy = vi.spyOn(auditService, "logIncomingConnectionCreated");

      await service.handleIncomingConnectionRequest(mockWorkspaceId, mockRequest);

      expect(auditSpy).toHaveBeenCalledWith({
        workspaceId: mockWorkspaceId,
        connectionId: mockConnection.id,
        remoteInstanceId: mockRequest.instanceId,
        remoteUrl: mockRequest.instanceUrl,
      });
    });

    it("should log connection rejected on invalid signature", async () => {
      vi.spyOn(signatureService, "verifyConnectionRequest").mockResolvedValue({
        valid: false,
        error: "Invalid signature",
      });
      const auditSpy = vi.spyOn(auditService, "logIncomingConnectionRejected");

      await expect(
        service.handleIncomingConnectionRequest(mockWorkspaceId, mockRequest)
      ).rejects.toThrow();

      expect(auditSpy).toHaveBeenCalledWith({
        workspaceId: mockWorkspaceId,
        remoteInstanceId: mockRequest.instanceId,
        remoteUrl: mockRequest.instanceUrl,
        reason: "Invalid signature",
        error: "Invalid signature",
      });
    });
  });
});