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

/**
 * Query Service Tests
 *
 * Tests for federated query message handling.
 */

import { describe, it, expect, beforeEach, vi } from "vitest";
import { Test, TestingModule } from "@nestjs/testing";
import { ConfigService } from "@nestjs/config";
import {
  FederationConnectionStatus,
  FederationMessageType,
  FederationMessageStatus,
} from "@prisma/client";
import { QueryService } from "./query.service";
import { PrismaService } from "../prisma/prisma.service";
import { FederationService } from "./federation.service";
import { SignatureService } from "./signature.service";
import { TasksService } from "../tasks/tasks.service";
import { EventsService } from "../events/events.service";
import { ProjectsService } from "../projects/projects.service";
import { HttpService } from "@nestjs/axios";
import { of, throwError } from "rxjs";
import type { AxiosResponse } from "axios";

describe("QueryService", () => {
  let service: QueryService;
  let prisma: PrismaService;
  let federationService: FederationService;
  let signatureService: SignatureService;
  let httpService: HttpService;

  const mockPrisma = {
    federationConnection: {
      findUnique: vi.fn(),
      findFirst: vi.fn(),
    },
    federationMessage: {
      create: vi.fn(),
      findMany: vi.fn(),
      findUnique: vi.fn(),
      findFirst: vi.fn(),
      update: vi.fn(),
    },
  };

  const mockFederationService = {
    getInstanceIdentity: vi.fn(),
    getPublicIdentity: vi.fn(),
  };

  const mockSignatureService = {
    signMessage: vi.fn(),
    verifyMessage: vi.fn(),
    validateTimestamp: vi.fn(),
  };

  const mockHttpService = {
    post: vi.fn(),
  };

  const mockConfig = {
    get: vi.fn(),
  };

  const mockTasksService = {
    findAll: vi.fn(),
  };

  const mockEventsService = {
    findAll: vi.fn(),
  };

  const mockProjectsService = {
    findAll: vi.fn(),
  };

  beforeEach(async () => {
    const module: TestingModule = await Test.createTestingModule({
      providers: [
        QueryService,
        { provide: PrismaService, useValue: mockPrisma },
        { provide: FederationService, useValue: mockFederationService },
        { provide: SignatureService, useValue: mockSignatureService },
        { provide: HttpService, useValue: mockHttpService },
        { provide: ConfigService, useValue: mockConfig },
        { provide: TasksService, useValue: mockTasksService },
        { provide: EventsService, useValue: mockEventsService },
        { provide: ProjectsService, useValue: mockProjectsService },
      ],
    }).compile();

    service = module.get<QueryService>(QueryService);
    prisma = module.get<PrismaService>(PrismaService);
    federationService = module.get<FederationService>(FederationService);
    signatureService = module.get<SignatureService>(SignatureService);
    httpService = module.get<HttpService>(HttpService);

    // Setup default mock return values for service queries
    mockTasksService.findAll.mockResolvedValue({
      data: [],
      meta: { total: 0, page: 1, limit: 50, totalPages: 0 },
    });
    mockEventsService.findAll.mockResolvedValue({
      data: [],
      meta: { total: 0, page: 1, limit: 50, totalPages: 0 },
    });
    mockProjectsService.findAll.mockResolvedValue({
      data: [],
      meta: { total: 0, page: 1, limit: 50, totalPages: 0 },
    });

    vi.clearAllMocks();
  });

  describe("sendQuery", () => {
    it("should send query to remote instance with signed message", async () => {
      const workspaceId = "workspace-1";
      const connectionId = "connection-1";
      const query = "SELECT * FROM tasks";
      const context = { userId: "user-1" };

      const mockConnection = {
        id: connectionId,
        workspaceId,
        remoteInstanceId: "remote-instance-1",
        remoteUrl: "https://remote.example.com",
        remotePublicKey: "mock-public-key",
        status: FederationConnectionStatus.ACTIVE,
      };

      const mockIdentity = {
        id: "identity-1",
        instanceId: "local-instance-1",
        name: "Local Instance",
        url: "https://local.example.com",
        publicKey: "local-public-key",
        privateKey: "local-private-key",
      };

      const mockMessage = {
        id: "message-1",
        workspaceId,
        connectionId,
        messageType: FederationMessageType.QUERY,
        messageId: expect.any(String),
        correlationId: null,
        query,
        response: null,
        status: FederationMessageStatus.PENDING,
        error: null,
        signature: "mock-signature",
        createdAt: new Date(),
        updatedAt: new Date(),
        deliveredAt: null,
      };

      const mockResponse: AxiosResponse = {
        data: { success: true },
        status: 200,
        statusText: "OK",
        headers: {},
        config: { headers: {} as never },
      };

      mockPrisma.federationConnection.findUnique.mockResolvedValue(mockConnection);
      mockFederationService.getInstanceIdentity.mockResolvedValue(mockIdentity);
      mockSignatureService.signMessage.mockResolvedValue("mock-signature");
      mockPrisma.federationMessage.create.mockResolvedValue(mockMessage);
      mockHttpService.post.mockReturnValue(of(mockResponse));

      const result = await service.sendQuery(workspaceId, connectionId, query, context);

      expect(result).toBeDefined();
      expect(result.messageType).toBe(FederationMessageType.QUERY);
      expect(result.query).toBe(query);
      expect(mockPrisma.federationConnection.findUnique).toHaveBeenCalledWith({
        where: {
          id: connectionId,
          workspaceId,
          status: FederationConnectionStatus.ACTIVE,
        },
      });
      expect(mockPrisma.federationMessage.create).toHaveBeenCalled();
      expect(mockHttpService.post).toHaveBeenCalledWith(
        `${mockConnection.remoteUrl}/api/v1/federation/incoming/query`,
        expect.objectContaining({
          messageId: expect.any(String),
          instanceId: mockIdentity.instanceId,
          query,
          context,
          timestamp: expect.any(Number),
          signature: "mock-signature",
        })
      );
    });

    it("should throw error if connection not found", async () => {
      mockPrisma.federationConnection.findUnique.mockResolvedValue(null);

      await expect(
        service.sendQuery("workspace-1", "connection-1", "SELECT * FROM tasks")
      ).rejects.toThrow("Connection not found");
    });

    it("should throw error if connection not active", async () => {
      // Connection should not be found by query because it's not ACTIVE
      mockPrisma.federationConnection.findUnique.mockResolvedValue(null);

      await expect(
        service.sendQuery("workspace-1", "connection-1", "SELECT * FROM tasks")
      ).rejects.toThrow("Connection not found");

      // Verify that findUnique was called with status: ACTIVE in the query
      expect(mockPrisma.federationConnection.findUnique).toHaveBeenCalledWith({
        where: {
          id: "connection-1",
          workspaceId: "workspace-1",
          status: FederationConnectionStatus.ACTIVE,
        },
      });
    });

    it("should handle network errors gracefully", async () => {
      const mockConnection = {
        id: "connection-1",
        workspaceId: "workspace-1",
        remoteInstanceId: "remote-instance-1",
        remoteUrl: "https://remote.example.com",
        status: FederationConnectionStatus.ACTIVE,
      };

      const mockIdentity = {
        instanceId: "local-instance-1",
      };

      mockPrisma.federationConnection.findUnique.mockResolvedValue(mockConnection);
      mockFederationService.getInstanceIdentity.mockResolvedValue(mockIdentity);
      mockSignatureService.signMessage.mockResolvedValue("mock-signature");
      mockPrisma.federationMessage.create.mockResolvedValue({
        id: "message-1",
        messageId: "msg-1",
      });
      mockHttpService.post.mockReturnValue(throwError(() => new Error("Network error")));

      await expect(
        service.sendQuery("workspace-1", "connection-1", "SELECT * FROM tasks")
      ).rejects.toThrow("Failed to send query");
    });
  });

  describe("handleIncomingQuery", () => {
    it("should process valid incoming query", async () => {
      const queryMessage = {
        messageId: "msg-1",
        instanceId: "remote-instance-1",
        query: "SELECT * FROM tasks",
        context: {},
        timestamp: Date.now(),
        signature: "valid-signature",
      };

      const mockConnection = {
        id: "connection-1",
        workspaceId: "workspace-1",
        remoteInstanceId: "remote-instance-1",
        status: FederationConnectionStatus.ACTIVE,
      };

      const mockIdentity = {
        instanceId: "local-instance-1",
      };

      mockPrisma.federationConnection.findFirst.mockResolvedValue(mockConnection);
      mockSignatureService.validateTimestamp.mockReturnValue(true);
      mockSignatureService.verifyMessage.mockResolvedValue({ valid: true });
      mockFederationService.getInstanceIdentity.mockResolvedValue(mockIdentity);
      mockSignatureService.signMessage.mockResolvedValue("response-signature");

      const result = await service.handleIncomingQuery(queryMessage);

      expect(result).toBeDefined();
      expect(result.messageId).toBeDefined();
      expect(result.correlationId).toBe(queryMessage.messageId);
      expect(result.instanceId).toBe(mockIdentity.instanceId);
      expect(result.signature).toBe("response-signature");
    });

    it("should reject query with invalid signature", async () => {
      const queryMessage = {
        messageId: "msg-1",
        instanceId: "remote-instance-1",
        query: "SELECT * FROM tasks",
        timestamp: Date.now(),
        signature: "invalid-signature",
      };

      const mockConnection = {
        id: "connection-1",
        workspaceId: "workspace-1",
        remoteInstanceId: "remote-instance-1",
        status: FederationConnectionStatus.ACTIVE,
      };

      mockPrisma.federationConnection.findFirst.mockResolvedValue(mockConnection);
      mockSignatureService.validateTimestamp.mockReturnValue(true);
      mockSignatureService.verifyMessage.mockResolvedValue({
        valid: false,
        error: "Invalid signature",
      });

      await expect(service.handleIncomingQuery(queryMessage)).rejects.toThrow("Invalid signature");
    });

    it("should reject query with expired timestamp", async () => {
      const queryMessage = {
        messageId: "msg-1",
        instanceId: "remote-instance-1",
        query: "SELECT * FROM tasks",
        timestamp: Date.now() - 10 * 60 * 1000, // 10 minutes ago
        signature: "valid-signature",
      };

      const mockConnection = {
        id: "connection-1",
        workspaceId: "workspace-1",
        remoteInstanceId: "remote-instance-1",
        status: FederationConnectionStatus.ACTIVE,
      };

      mockPrisma.federationConnection.findFirst.mockResolvedValue(mockConnection);
      mockSignatureService.validateTimestamp.mockReturnValue(false);

      await expect(service.handleIncomingQuery(queryMessage)).rejects.toThrow(
        "Query timestamp is outside acceptable range"
      );
    });

    it("should reject query from inactive connection", async () => {
      const queryMessage = {
        messageId: "msg-1",
        instanceId: "remote-instance-1",
        query: "SELECT * FROM tasks",
        timestamp: Date.now(),
        signature: "valid-signature",
      };

      // Connection should not be found because status filter in query excludes non-ACTIVE
      mockPrisma.federationConnection.findFirst.mockResolvedValue(null);
      mockSignatureService.validateTimestamp.mockReturnValue(true);

      await expect(service.handleIncomingQuery(queryMessage)).rejects.toThrow(
        "No connection found for remote instance"
      );

      // Verify the findFirst was called with status: ACTIVE filter
      expect(mockPrisma.federationConnection.findFirst).toHaveBeenCalledWith({
        where: {
          remoteInstanceId: "remote-instance-1",
          status: FederationConnectionStatus.ACTIVE,
        },
      });
    });

    it("should reject query from unknown instance", async () => {
      const queryMessage = {
        messageId: "msg-1",
        instanceId: "unknown-instance",
        query: "SELECT * FROM tasks",
        timestamp: Date.now(),
        signature: "valid-signature",
      };

      mockPrisma.federationConnection.findFirst.mockResolvedValue(null);
      mockSignatureService.validateTimestamp.mockReturnValue(true);

      await expect(service.handleIncomingQuery(queryMessage)).rejects.toThrow(
        "No connection found for remote instance"
      );
    });
  });

  describe("getQueryMessages", () => {
    it("should return query messages for workspace", async () => {
      const workspaceId = "workspace-1";
      const mockMessages = [
        {
          id: "msg-1",
          workspaceId,
          connectionId: "connection-1",
          messageType: FederationMessageType.QUERY,
          messageId: "unique-msg-1",
          query: "SELECT * FROM tasks",
          status: FederationMessageStatus.DELIVERED,
          createdAt: new Date(),
          updatedAt: new Date(),
        },
      ];

      mockPrisma.federationMessage.findMany.mockResolvedValue(mockMessages);

      const result = await service.getQueryMessages(workspaceId);

      expect(result).toHaveLength(1);
      expect(result[0].id).toBe("msg-1");
      expect(mockPrisma.federationMessage.findMany).toHaveBeenCalledWith({
        where: {
          workspaceId,
          messageType: FederationMessageType.QUERY,
        },
        orderBy: { createdAt: "desc" },
      });
    });

    it("should filter by status when provided", async () => {
      const workspaceId = "workspace-1";
      const status = FederationMessageStatus.PENDING;

      mockPrisma.federationMessage.findMany.mockResolvedValue([]);

      await service.getQueryMessages(workspaceId, status);

      expect(mockPrisma.federationMessage.findMany).toHaveBeenCalledWith({
        where: {
          workspaceId,
          messageType: FederationMessageType.QUERY,
          status,
        },
        orderBy: { createdAt: "desc" },
      });
    });
  });

  describe("getQueryMessage", () => {
    it("should return query message by ID", async () => {
      const workspaceId = "workspace-1";
      const messageId = "msg-1";
      const mockMessage = {
        id: "msg-1",
        workspaceId,
        messageType: FederationMessageType.QUERY,
        messageId: "unique-msg-1",
        query: "SELECT * FROM tasks",
        status: FederationMessageStatus.DELIVERED,
        createdAt: new Date(),
        updatedAt: new Date(),
      };

      mockPrisma.federationMessage.findUnique.mockResolvedValue(mockMessage);

      const result = await service.getQueryMessage(workspaceId, messageId);

      expect(result).toBeDefined();
      expect(result.id).toBe(messageId);
      expect(mockPrisma.federationMessage.findUnique).toHaveBeenCalledWith({
        where: { id: messageId, workspaceId },
      });
    });

    it("should throw error if message not found", async () => {
      mockPrisma.federationMessage.findUnique.mockResolvedValue(null);

      await expect(service.getQueryMessage("workspace-1", "msg-1")).rejects.toThrow(
        "Query message not found"
      );
    });
  });

  describe("processQueryResponse", () => {
    it("should update message with response", async () => {
      const response = {
        messageId: "response-1",
        correlationId: "original-msg-1",
        instanceId: "remote-instance-1",
        success: true,
        data: { tasks: [] },
        timestamp: Date.now(),
        signature: "valid-signature",
      };

      const mockMessage = {
        id: "msg-1",
        messageId: "original-msg-1",
        workspaceId: "workspace-1",
        status: FederationMessageStatus.PENDING,
      };

      mockPrisma.federationMessage.findFirst.mockResolvedValue(mockMessage);
      mockSignatureService.validateTimestamp.mockReturnValue(true);
      mockSignatureService.verifyMessage.mockResolvedValue({ valid: true });
      mockPrisma.federationMessage.update.mockResolvedValue({
        ...mockMessage,
        status: FederationMessageStatus.DELIVERED,
        response: response.data,
      });

      await service.processQueryResponse(response);

      expect(mockPrisma.federationMessage.update).toHaveBeenCalledWith({
        where: { id: mockMessage.id },
        data: {
          status: FederationMessageStatus.DELIVERED,
          response: response.data,
          deliveredAt: expect.any(Date),
        },
      });
    });

    it("should reject response with invalid signature", async () => {
      const response = {
        messageId: "response-1",
        correlationId: "original-msg-1",
        instanceId: "remote-instance-1",
        success: true,
        timestamp: Date.now(),
        signature: "invalid-signature",
      };

      const mockMessage = {
        id: "msg-1",
        messageId: "original-msg-1",
        workspaceId: "workspace-1",
      };

      mockPrisma.federationMessage.findFirst.mockResolvedValue(mockMessage);
      mockSignatureService.validateTimestamp.mockReturnValue(true);
      mockSignatureService.verifyMessage.mockResolvedValue({
        valid: false,
        error: "Invalid signature",
      });

      await expect(service.processQueryResponse(response)).rejects.toThrow("Invalid signature");
    });
  });

  describe("query processing with actual data", () => {
    it("should process 'get tasks' query and return tasks data", async () => {
      const queryMessage = {
        messageId: "msg-1",
        instanceId: "remote-instance-1",
        query: "get tasks",
        context: { workspaceId: "workspace-1" },
        timestamp: Date.now(),
        signature: "valid-signature",
      };

      const mockConnection = {
        id: "connection-1",
        workspaceId: "workspace-1",
        remoteInstanceId: "remote-instance-1",
        status: FederationConnectionStatus.ACTIVE,
      };

      const mockIdentity = {
        instanceId: "local-instance-1",
      };

      mockPrisma.federationConnection.findFirst.mockResolvedValue(mockConnection);
      mockSignatureService.validateTimestamp.mockReturnValue(true);
      mockSignatureService.verifyMessage.mockResolvedValue({ valid: true });
      mockFederationService.getInstanceIdentity.mockResolvedValue(mockIdentity);
      mockSignatureService.signMessage.mockResolvedValue("response-signature");

      const result = await service.handleIncomingQuery(queryMessage);

      expect(result).toBeDefined();
      expect(result.success).toBe(true);
      expect(result.data).toBeDefined();
      expect(result.correlationId).toBe(queryMessage.messageId);
    });

    it("should process 'get events' query and return events data", async () => {
      const queryMessage = {
        messageId: "msg-2",
        instanceId: "remote-instance-1",
        query: "get events",
        context: { workspaceId: "workspace-1" },
        timestamp: Date.now(),
        signature: "valid-signature",
      };

      const mockConnection = {
        id: "connection-1",
        workspaceId: "workspace-1",
        remoteInstanceId: "remote-instance-1",
        status: FederationConnectionStatus.ACTIVE,
      };

      const mockIdentity = {
        instanceId: "local-instance-1",
      };

      mockPrisma.federationConnection.findFirst.mockResolvedValue(mockConnection);
      mockSignatureService.validateTimestamp.mockReturnValue(true);
      mockSignatureService.verifyMessage.mockResolvedValue({ valid: true });
      mockFederationService.getInstanceIdentity.mockResolvedValue(mockIdentity);
      mockSignatureService.signMessage.mockResolvedValue("response-signature");

      const result = await service.handleIncomingQuery(queryMessage);

      expect(result).toBeDefined();
      expect(result.success).toBe(true);
      expect(result.data).toBeDefined();
    });

    it("should process 'get projects' query and return projects data", async () => {
      const queryMessage = {
        messageId: "msg-3",
        instanceId: "remote-instance-1",
        query: "get projects",
        context: { workspaceId: "workspace-1" },
        timestamp: Date.now(),
        signature: "valid-signature",
      };

      const mockConnection = {
        id: "connection-1",
        workspaceId: "workspace-1",
        remoteInstanceId: "remote-instance-1",
        status: FederationConnectionStatus.ACTIVE,
      };

      const mockIdentity = {
        instanceId: "local-instance-1",
      };

      mockPrisma.federationConnection.findFirst.mockResolvedValue(mockConnection);
      mockSignatureService.validateTimestamp.mockReturnValue(true);
      mockSignatureService.verifyMessage.mockResolvedValue({ valid: true });
      mockFederationService.getInstanceIdentity.mockResolvedValue(mockIdentity);
      mockSignatureService.signMessage.mockResolvedValue("response-signature");

      const result = await service.handleIncomingQuery(queryMessage);

      expect(result).toBeDefined();
      expect(result.success).toBe(true);
      expect(result.data).toBeDefined();
    });

    it("should handle unknown query type gracefully", async () => {
      const queryMessage = {
        messageId: "msg-4",
        instanceId: "remote-instance-1",
        query: "unknown command",
        context: { workspaceId: "workspace-1" },
        timestamp: Date.now(),
        signature: "valid-signature",
      };

      const mockConnection = {
        id: "connection-1",
        workspaceId: "workspace-1",
        remoteInstanceId: "remote-instance-1",
        status: FederationConnectionStatus.ACTIVE,
      };

      const mockIdentity = {
        instanceId: "local-instance-1",
      };

      mockPrisma.federationConnection.findFirst.mockResolvedValue(mockConnection);
      mockSignatureService.validateTimestamp.mockReturnValue(true);
      mockSignatureService.verifyMessage.mockResolvedValue({ valid: true });
      mockFederationService.getInstanceIdentity.mockResolvedValue(mockIdentity);
      mockSignatureService.signMessage.mockResolvedValue("response-signature");

      const result = await service.handleIncomingQuery(queryMessage);

      expect(result).toBeDefined();
      expect(result.success).toBe(false);
      expect(result.error).toContain("Unknown query type");
    });

    it("should enforce workspace context in queries", async () => {
      const queryMessage = {
        messageId: "msg-5",
        instanceId: "remote-instance-1",
        query: "get tasks",
        context: {}, // Missing workspaceId
        timestamp: Date.now(),
        signature: "valid-signature",
      };

      const mockConnection = {
        id: "connection-1",
        workspaceId: "workspace-1",
        remoteInstanceId: "remote-instance-1",
        status: FederationConnectionStatus.ACTIVE,
      };

      const mockIdentity = {
        instanceId: "local-instance-1",
      };

      mockPrisma.federationConnection.findFirst.mockResolvedValue(mockConnection);
      mockSignatureService.validateTimestamp.mockReturnValue(true);
      mockSignatureService.verifyMessage.mockResolvedValue({ valid: true });
      mockFederationService.getInstanceIdentity.mockResolvedValue(mockIdentity);
      mockSignatureService.signMessage.mockResolvedValue("response-signature");

      const result = await service.handleIncomingQuery(queryMessage);

      expect(result).toBeDefined();
      expect(result.success).toBe(false);
      expect(result.error).toContain("workspaceId");
    });

    it("should reject queries for UserCredential entity type", async () => {
      const queryMessage = {
        messageId: "msg-1",
        instanceId: "remote-instance-1",
        query: "credentials",
        context: { workspaceId: "workspace-1" },
        timestamp: Date.now(),
        signature: "valid-signature",
      };

      const mockConnection = {
        id: "connection-1",
        workspaceId: "workspace-1",
        remoteInstanceId: "remote-instance-1",
        status: FederationConnectionStatus.ACTIVE,
      };

      const mockIdentity = {
        instanceId: "local-instance-1",
      };

      mockPrisma.federationConnection.findFirst.mockResolvedValue(mockConnection);
      mockSignatureService.validateTimestamp.mockReturnValue(true);
      mockSignatureService.verifyMessage.mockResolvedValue({ valid: true });
      mockFederationService.getInstanceIdentity.mockResolvedValue(mockIdentity);
      mockSignatureService.signMessage.mockResolvedValue("response-signature");

      const result = await service.handleIncomingQuery(queryMessage);

      expect(result).toBeDefined();
      expect(result.success).toBe(false);
      expect(result.error).toContain("Credential queries are not allowed");
    });

    it("should reject queries containing credential-related keywords", async () => {
      const credentialQueries = [
        "SELECT * FROM user_credentials",
        "get all credentials",
        "show my api keys",
        "list oauth tokens",
      ];

      const mockConnection = {
        id: "connection-1",
        workspaceId: "workspace-1",
        remoteInstanceId: "remote-instance-1",
        status: FederationConnectionStatus.ACTIVE,
      };

      const mockIdentity = {
        instanceId: "local-instance-1",
      };

      mockPrisma.federationConnection.findFirst.mockResolvedValue(mockConnection);
      mockSignatureService.validateTimestamp.mockReturnValue(true);
      mockSignatureService.verifyMessage.mockResolvedValue({ valid: true });
      mockFederationService.getInstanceIdentity.mockResolvedValue(mockIdentity);
      mockSignatureService.signMessage.mockResolvedValue("response-signature");

      for (const query of credentialQueries) {
        const queryMessage = {
          messageId: `msg-${Math.random()}`,
          instanceId: "remote-instance-1",
          query,
          context: { workspaceId: "workspace-1" },
          timestamp: Date.now(),
          signature: "valid-signature",
        };

        const result = await service.handleIncomingQuery(queryMessage);

        expect(result.success).toBe(false);
        expect(result.error).toContain("Credential queries are not allowed");
      }
    });
  });
});