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

/**
 * Query Controller Tests
 *
 * Tests for federated query API endpoints.
 */

import { describe, it, expect, beforeEach, vi } from "vitest";
import { Test, TestingModule } from "@nestjs/testing";
import { FederationMessageType, FederationMessageStatus } from "@prisma/client";
import { QueryController } from "./query.controller";
import { QueryService } from "./query.service";
import { AuthGuard } from "../auth/guards/auth.guard";
import type { AuthenticatedRequest } from "../common/types/user.types";
import type { SendQueryDto, IncomingQueryDto } from "./dto/query.dto";

describe("QueryController", () => {
  let controller: QueryController;
  let queryService: QueryService;

  const mockQueryService = {
    sendQuery: vi.fn(),
    handleIncomingQuery: vi.fn(),
    getQueryMessages: vi.fn(),
    getQueryMessage: vi.fn(),
  };

  beforeEach(async () => {
    const module: TestingModule = await Test.createTestingModule({
      controllers: [QueryController],
      providers: [{ provide: QueryService, useValue: mockQueryService }],
    })
      .overrideGuard(AuthGuard)
      .useValue({ canActivate: () => true })
      .compile();

    controller = module.get<QueryController>(QueryController);
    queryService = module.get<QueryService>(QueryService);

    vi.clearAllMocks();
  });

  describe("sendQuery", () => {
    it("should send query to remote instance", async () => {
      const req = {
        user: {
          id: "user-1",
          workspaceId: "workspace-1",
        },
      } as AuthenticatedRequest;

      const dto: SendQueryDto = {
        connectionId: "connection-1",
        query: "SELECT * FROM tasks",
        context: { userId: "user-1" },
      };

      const mockResult = {
        id: "msg-1",
        workspaceId: "workspace-1",
        connectionId: "connection-1",
        messageType: FederationMessageType.QUERY,
        messageId: "unique-msg-1",
        query: dto.query,
        status: FederationMessageStatus.PENDING,
        createdAt: new Date(),
        updatedAt: new Date(),
      };

      mockQueryService.sendQuery.mockResolvedValue(mockResult);

      const result = await controller.sendQuery(req, dto);

      expect(result).toBeDefined();
      expect(result.messageType).toBe(FederationMessageType.QUERY);
      expect(mockQueryService.sendQuery).toHaveBeenCalledWith(
        "workspace-1",
        dto.connectionId,
        dto.query,
        dto.context
      );
    });

    it("should throw error if user not authenticated", async () => {
      const req = {} as AuthenticatedRequest;

      const dto: SendQueryDto = {
        connectionId: "connection-1",
        query: "SELECT * FROM tasks",
      };

      await expect(controller.sendQuery(req, dto)).rejects.toThrow(
        "Workspace ID not found in request"
      );
    });
  });

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

      const mockResponse = {
        messageId: "response-1",
        correlationId: dto.messageId,
        instanceId: "local-instance-1",
        success: true,
        data: { tasks: [] },
        timestamp: Date.now(),
        signature: "response-signature",
      };

      mockQueryService.handleIncomingQuery.mockResolvedValue(mockResponse);

      const result = await controller.handleIncomingQuery(dto);

      expect(result).toBeDefined();
      expect(result.correlationId).toBe(dto.messageId);
      expect(mockQueryService.handleIncomingQuery).toHaveBeenCalledWith(dto);
    });

    it("should return error response for invalid query", async () => {
      const dto: IncomingQueryDto = {
        messageId: "msg-1",
        instanceId: "remote-instance-1",
        query: "SELECT * FROM tasks",
        timestamp: Date.now(),
        signature: "invalid-signature",
      };

      mockQueryService.handleIncomingQuery.mockRejectedValue(new Error("Invalid signature"));

      await expect(controller.handleIncomingQuery(dto)).rejects.toThrow("Invalid signature");
    });
  });

  describe("getQueries", () => {
    it("should return query messages for workspace", async () => {
      const req = {
        user: {
          id: "user-1",
          workspaceId: "workspace-1",
        },
      } as AuthenticatedRequest;

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

      mockQueryService.getQueryMessages.mockResolvedValue(mockMessages);

      const result = await controller.getQueries(req, undefined);

      expect(result).toHaveLength(1);
      expect(mockQueryService.getQueryMessages).toHaveBeenCalledWith("workspace-1", undefined);
    });

    it("should filter by status when provided", async () => {
      const req = {
        user: {
          id: "user-1",
          workspaceId: "workspace-1",
        },
      } as AuthenticatedRequest;

      const status = FederationMessageStatus.PENDING;

      mockQueryService.getQueryMessages.mockResolvedValue([]);

      await controller.getQueries(req, status);

      expect(mockQueryService.getQueryMessages).toHaveBeenCalledWith("workspace-1", status);
    });

    it("should throw error if user not authenticated", async () => {
      const req = {} as AuthenticatedRequest;

      await expect(controller.getQueries(req, undefined)).rejects.toThrow(
        "Workspace ID not found in request"
      );
    });
  });

  describe("getQuery", () => {
    it("should return query message by ID", async () => {
      const req = {
        user: {
          id: "user-1",
          workspaceId: "workspace-1",
        },
      } as AuthenticatedRequest;

      const messageId = "msg-1";

      const mockMessage = {
        id: messageId,
        workspaceId: "workspace-1",
        connectionId: "connection-1",
        messageType: FederationMessageType.QUERY,
        messageId: "unique-msg-1",
        query: "SELECT * FROM tasks",
        status: FederationMessageStatus.DELIVERED,
        createdAt: new Date(),
        updatedAt: new Date(),
      };

      mockQueryService.getQueryMessage.mockResolvedValue(mockMessage);

      const result = await controller.getQuery(req, messageId);

      expect(result).toBeDefined();
      expect(result.id).toBe(messageId);
      expect(mockQueryService.getQueryMessage).toHaveBeenCalledWith("workspace-1", messageId);
    });

    it("should throw error if user not authenticated", async () => {
      const req = {} as AuthenticatedRequest;

      await expect(controller.getQuery(req, "msg-1")).rejects.toThrow(
        "Workspace ID not found in request"
      );
    });
  });
});