import { describe, it, expect, vi, beforeEach, afterEach } from "vitest"; import { SourceService } from "@/src/services/source-service"; import type { SourceType } from "@/src/types/source"; const queryMock = vi.fn(); const closeMock = vi.fn().mockResolvedValue(undefined); // Hoisted by Vitest above the static imports below, so // source-query-service.ts (which statically imports FabricOneLakeAdapter) // sees this mock too — no real mssql/AAD network calls happen in this file. vi.mock("@/src/adapters/fabric-onelake-adapter", () => { return { // Must be a real `function` (not an arrow fn) so `new FabricOneLakeAdapter(...)` // in source-query-service.ts works: a constructor that explicitly returns an // object substitutes that object for `this`, per JS semantics. FabricOneLakeAdapter: vi.fn().mockImplementation(function FakeFabricOneLakeAdapter() { return { query: queryMock, close: closeMock }; }), }; }); import { SourceQueryService, SourceQueryTimeoutError, SourceQueryUnsupportedError, SOURCE_QUERY_HARD_MAX_ROWS, SOURCE_QUERY_HARD_MAX_TIMEOUT_MS, } from "@/src/services/source-query-service"; const BASE_CONNECTION_CONFIG = { credentials: { tenantId: "t", clientId: "c", clientSecret: "s" }, sqlEndpoint: "ep.datawarehouse.fabric.microsoft.com", database: "lh_test", }; async function registerSource( sourceService: SourceService, opts: { source_id: string; adapter_id?: string; source_type?: SourceType }, ): Promise { await sourceService.createSource({ source_id: opts.source_id, tenant_id: "t1", project_id: "p1", source_type: opts.source_type ?? "database", name: "Fabric Warehouse", adapter_id: opts.adapter_id ?? "fabric_onelake_sql", }); return opts.source_id; } describe("SourceQueryService.querySource — bounded execution", () => { beforeEach(() => { queryMock.mockReset(); closeMock.mockClear(); }); it("executes a query against a registered fabric_onelake_sql database source and returns rows", async () => { const sourceService = new SourceService(); const sourceId = await registerSource(sourceService, { source_id: "src_query_ok" }); queryMock.mockResolvedValue({ rows: [{ id: 1 }, { id: 2 }], rowCount: 2, durationMs: 5 }); const service = new SourceQueryService({ sourceService }); const result = await service.querySource({ source_id: sourceId, query: "SELECT * FROM dbo.customers", tenant_id: "t1", project_id: "p1", connection_config: BASE_CONNECTION_CONFIG, }); expect(result.rows).toEqual([{ id: 1 }, { id: 2 }]); expect(result.row_count).toBe(2); expect(result.truncated).toBe(false); expect(result.bounds.max_rows).toBeGreaterThan(0); expect(closeMock).toHaveBeenCalledTimes(1); }); it("rejects a non-database source with SourceQueryUnsupportedError", async () => { const sourceService = new SourceService(); const sourceId = await registerSource(sourceService, { source_id: "src_file_1", source_type: "file", adapter_id: undefined }); const service = new SourceQueryService({ sourceService }); await expect( service.querySource({ source_id: sourceId, query: "SELECT 1", tenant_id: "t1", project_id: "p1", connection_config: BASE_CONNECTION_CONFIG }), ).rejects.toBeInstanceOf(SourceQueryUnsupportedError); }); it("rejects a database source with an unsupported adapter_id", async () => { const sourceService = new SourceService(); const sourceId = await registerSource(sourceService, { source_id: "src_bad_adapter", adapter_id: "some_other_adapter" }); const service = new SourceQueryService({ sourceService }); await expect( service.querySource({ source_id: sourceId, query: "SELECT 1", tenant_id: "t1", project_id: "p1", connection_config: BASE_CONNECTION_CONFIG }), ).rejects.toBeInstanceOf(SourceQueryUnsupportedError); }); it("propagates a 'not found' error for an unregistered source, without ever building an adapter", async () => { const sourceService = new SourceService(); const service = new SourceQueryService({ sourceService }); await expect( service.querySource({ source_id: "does_not_exist", query: "SELECT 1", tenant_id: "t1", project_id: "p1", connection_config: BASE_CONNECTION_CONFIG }), ).rejects.toThrow(/not found/i); expect(queryMock).not.toHaveBeenCalled(); }); it("truncates rows to max_rows and flags truncated:true", async () => { const sourceService = new SourceService(); const sourceId = await registerSource(sourceService, { source_id: "src_truncate" }); const rows = Array.from({ length: 50 }, (_, i) => ({ id: i })); queryMock.mockResolvedValue({ rows, rowCount: rows.length, durationMs: 3 }); const service = new SourceQueryService({ sourceService }); const result = await service.querySource({ source_id: sourceId, query: "SELECT * FROM dbo.big_table", tenant_id: "t1", project_id: "p1", max_rows: 10, connection_config: BASE_CONNECTION_CONFIG, }); expect(result.rows).toHaveLength(10); expect(result.row_count).toBe(10); expect(result.truncated).toBe(true); expect(result.bounds.max_rows).toBe(10); }); it("does not flag truncated when row count is exactly max_rows", async () => { const sourceService = new SourceService(); const sourceId = await registerSource(sourceService, { source_id: "src_exact" }); const rows = Array.from({ length: 10 }, (_, i) => ({ id: i })); queryMock.mockResolvedValue({ rows, rowCount: rows.length, durationMs: 1 }); const service = new SourceQueryService({ sourceService }); const result = await service.querySource({ source_id: sourceId, query: "SELECT * FROM dbo.t", tenant_id: "t1", project_id: "p1", max_rows: 10, connection_config: BASE_CONNECTION_CONFIG, }); expect(result.truncated).toBe(false); expect(result.rows).toHaveLength(10); }); it("clamps a caller-requested max_rows above the hard ceiling instead of allowing an unbounded read", async () => { const sourceService = new SourceService(); const sourceId = await registerSource(sourceService, { source_id: "src_clamp_rows" }); queryMock.mockResolvedValue({ rows: [{ id: 1 }], rowCount: 1, durationMs: 1 }); const service = new SourceQueryService({ sourceService }); const result = await service.querySource({ source_id: sourceId, query: "SELECT 1", tenant_id: "t1", project_id: "p1", max_rows: SOURCE_QUERY_HARD_MAX_ROWS * 10, connection_config: BASE_CONNECTION_CONFIG, }); expect(result.bounds.max_rows).toBe(SOURCE_QUERY_HARD_MAX_ROWS); }); it("clamps a caller-requested timeout_ms above the hard ceiling", async () => { const sourceService = new SourceService(); const sourceId = await registerSource(sourceService, { source_id: "src_clamp_timeout" }); queryMock.mockResolvedValue({ rows: [], rowCount: 0, durationMs: 1 }); const service = new SourceQueryService({ sourceService }); const result = await service.querySource({ source_id: sourceId, query: "SELECT 1", tenant_id: "t1", project_id: "p1", timeout_ms: SOURCE_QUERY_HARD_MAX_TIMEOUT_MS * 100, connection_config: BASE_CONNECTION_CONFIG, }); expect(result.bounds.timeout_ms).toBe(SOURCE_QUERY_HARD_MAX_TIMEOUT_MS); }); it("times out a query that runs longer than the bounded timeout, with SourceQueryTimeoutError", async () => { const sourceService = new SourceService(); const sourceId = await registerSource(sourceService, { source_id: "src_timeout" }); queryMock.mockImplementation( () => new Promise((resolve) => setTimeout(() => resolve({ rows: [], rowCount: 0, durationMs: 1 }), 300)), ); const service = new SourceQueryService({ sourceService }); await expect( service.querySource({ source_id: sourceId, query: "SELECT * FROM dbo.slow_view", tenant_id: "t1", project_id: "p1", timeout_ms: 20, connection_config: BASE_CONNECTION_CONFIG, }), ).rejects.toBeInstanceOf(SourceQueryTimeoutError); }); it("still releases the adapter (calls close) after a timeout", async () => { const sourceService = new SourceService(); const sourceId = await registerSource(sourceService, { source_id: "src_timeout_cleanup" }); queryMock.mockImplementation( () => new Promise((resolve) => setTimeout(() => resolve({ rows: [], rowCount: 0, durationMs: 1 }), 300)), ); const service = new SourceQueryService({ sourceService }); await expect( service.querySource({ source_id: sourceId, query: "SELECT 1", tenant_id: "t1", project_id: "p1", timeout_ms: 20, connection_config: BASE_CONNECTION_CONFIG, }), ).rejects.toThrow(); expect(closeMock).toHaveBeenCalledTimes(1); }); it("propagates adapter query errors as-is", async () => { const sourceService = new SourceService(); const sourceId = await registerSource(sourceService, { source_id: "src_adapter_error" }); queryMock.mockRejectedValue(new Error("Invalid object name 'dbo.does_not_exist'")); const service = new SourceQueryService({ sourceService }); await expect( service.querySource({ source_id: sourceId, query: "SELECT * FROM dbo.does_not_exist", tenant_id: "t1", project_id: "p1", connection_config: BASE_CONNECTION_CONFIG, }), ).rejects.toThrow(/Invalid object name/); }); }); describe("SourceQueryService — AUD0 0pt.source_queried event emission", () => { const originalFetch = global.fetch; afterEach(() => { global.fetch = originalFetch; }); it("emits a 0pt.source_queried AUD0 event with ok:true on success", async () => { queryMock.mockReset(); closeMock.mockClear(); const sourceService = new SourceService(); const sourceId = await registerSource(sourceService, { source_id: "src_aud0_ok" }); queryMock.mockResolvedValue({ rows: [{ id: 1 }], rowCount: 1, durationMs: 2 }); const calls: Array<{ url: string; body: Record }> = []; global.fetch = vi.fn(async (url: unknown, init: unknown) => { calls.push({ url: String(url), body: JSON.parse(String((init as RequestInit).body)) }); return new Response("{}", { status: 200 }); }) as unknown as typeof fetch; const service = new SourceQueryService({ sourceService, aud0BaseUrl: "http://aud0.test" }); await service.querySource({ source_id: sourceId, query: "SELECT * FROM dbo.customers", tenant_id: "t1", project_id: "p1", actor_id: "user_42", correlation_id: "corr_abc", connection_config: BASE_CONNECTION_CONFIG, }); expect(calls).toHaveLength(1); expect(calls[0].url).toBe("http://aud0.test/v1/aud0/events"); expect(calls[0].body.event_type).toBe("0pt.source_queried"); expect(calls[0].body.capability_id).toBe("0pt.query_source"); expect(calls[0].body.app_id).toBe("0pt"); expect(calls[0].body.resource_type).toBe("source"); expect(calls[0].body.resource_id).toBe(sourceId); expect(calls[0].body.actor_id).toBe("user_42"); expect(calls[0].body.correlation_id).toBe("corr_abc"); const payload = calls[0].body.payload as Record; expect(payload.ok).toBe(true); expect(payload.row_count).toBe(1); expect(payload.adapter_id).toBe("fabric_onelake_sql"); }); it("emits a 0pt.source_queried AUD0 event with ok:false and the error message on failure", async () => { queryMock.mockReset(); closeMock.mockClear(); const sourceService = new SourceService(); const sourceId = await registerSource(sourceService, { source_id: "src_aud0_fail" }); queryMock.mockRejectedValue(new Error("boom")); const calls: Array<{ url: string; body: Record }> = []; global.fetch = vi.fn(async (url: unknown, init: unknown) => { calls.push({ url: String(url), body: JSON.parse(String((init as RequestInit).body)) }); return new Response("{}", { status: 200 }); }) as unknown as typeof fetch; const service = new SourceQueryService({ sourceService, aud0BaseUrl: "http://aud0.test" }); await expect( service.querySource({ source_id: sourceId, query: "SELECT 1", tenant_id: "t1", project_id: "p1", connection_config: BASE_CONNECTION_CONFIG }), ).rejects.toThrow("boom"); expect(calls).toHaveLength(1); const payload = calls[0].body.payload as Record; expect(payload.ok).toBe(false); expect(payload.error).toBe("boom"); }); });