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