@neverinfamous/mysql-mcp 2.1.0 → 2.3.0

package/src/adapters/mysql/tools/security/__tests__/audit.test.ts

@@ -1,181 +1,239 @@
 /**
  * mysql-mcp - Security Audit Tools Unit Tests
- * 
+ *
  * Comprehensive tests for audit.ts module functions.
  */
 
-import { describe, it, expect, vi, beforeEach } from 'vitest';
+import { describe, it, expect, vi, beforeEach } from "vitest";
 import {
-    createSecurityAuditTool,
-    createSecurityFirewallStatusTool,
-    createSecurityFirewallRulesTool
-} from '../audit.js';
-import type { MySQLAdapter } from '../../../MySQLAdapter.js';
-import { createMockMySQLAdapter, createMockRequestContext, createMockQueryResult } from '../../../../../__tests__/mocks/index.js';
-
-describe('Security Audit Tools', () => {
-    let mockAdapter: ReturnType<typeof createMockMySQLAdapter>;
-    let mockContext: ReturnType<typeof createMockRequestContext>;
-
-    beforeEach(() => {
-        vi.clearAllMocks();
-        mockAdapter = createMockMySQLAdapter();
-        mockContext = createMockRequestContext();
+  createSecurityAuditTool,
+  createSecurityFirewallStatusTool,
+  createSecurityFirewallRulesTool,
+} from "../audit.js";
+import type { MySQLAdapter } from "../../../MySQLAdapter.js";
+import {
+  createMockMySQLAdapter,
+  createMockRequestContext,
+  createMockQueryResult,
+} from "../../../../../__tests__/mocks/index.js";
+
+describe("Security Audit Tools", () => {
+  let mockAdapter: ReturnType<typeof createMockMySQLAdapter>;
+  let mockContext: ReturnType<typeof createMockRequestContext>;
+
+  beforeEach(() => {
+    vi.clearAllMocks();
+    mockAdapter = createMockMySQLAdapter();
+    mockContext = createMockRequestContext();
+  });
+
+  describe("createSecurityAuditTool", () => {
+    it("should query mysql.audit_log if available", async () => {
+      // Mock table check
+      mockAdapter.executeQuery.mockResolvedValueOnce(
+        createMockQueryResult([{ TABLE_NAME: "audit_log" }]),
+      );
+
+      // Mock log query
+      mockAdapter.executeQuery.mockResolvedValueOnce(
+        createMockQueryResult([{ user: "test", event_type: "LOGIN" }]),
+      );
+
+      const tool = createSecurityAuditTool(
+        mockAdapter as unknown as MySQLAdapter,
+      );
+      const result = (await tool.handler(
+        {
+          limit: 10,
+          user: "test",
+        },
+        mockContext,
+      )) as { source: string; events: any[] };
+
+      expect(mockAdapter.executeQuery).toHaveBeenCalledTimes(2);
+      expect(result.source).toBe("mysql.audit_log");
+      expect(result.events).toHaveLength(1);
+
+      const queryCall = mockAdapter.executeQuery.mock.calls[1][0] as string;
+      expect(queryCall).toContain("FROM mysql.audit_log");
+      expect(queryCall).toContain("WHERE user LIKE ?");
+      expect(queryCall).toContain("WHERE user LIKE ?");
     });
 
-    describe('createSecurityAuditTool', () => {
-        it('should query mysql.audit_log if available', async () => {
-            // Mock table check
-            mockAdapter.executeQuery.mockResolvedValueOnce(createMockQueryResult([{ TABLE_NAME: 'audit_log' }]));
-
-            // Mock log query
-            mockAdapter.executeQuery.mockResolvedValueOnce(createMockQueryResult([
-                { user: 'test', event_type: 'LOGIN' }
-            ]));
-
-            const tool = createSecurityAuditTool(mockAdapter as unknown as MySQLAdapter);
-            const result = await tool.handler({
-                limit: 10,
-                user: 'test'
-            }, mockContext) as { source: string, events: any[] };
-
-            expect(mockAdapter.executeQuery).toHaveBeenCalledTimes(2);
-            expect(result.source).toBe('mysql.audit_log');
-            expect(result.events).toHaveLength(1);
-
-            const queryCall = mockAdapter.executeQuery.mock.calls[1][0] as string;
-            expect(queryCall).toContain('FROM mysql.audit_log');
-            expect(queryCall).toContain('WHERE user LIKE ?');
-            expect(queryCall).toContain('WHERE user LIKE ?');
-        });
-
-        it('should apply all filters in normal mode', async () => {
-            mockAdapter.executeQuery.mockResolvedValueOnce(createMockQueryResult([{ TABLE_NAME: 'audit_log' }]));
-            mockAdapter.executeQuery.mockResolvedValueOnce(createMockQueryResult([]));
-
-            const tool = createSecurityAuditTool(mockAdapter as unknown as MySQLAdapter);
-            await tool.handler({
-                limit: 10,
-                user: 'test',
-                eventType: 'LOGIN',
-                startTime: '2023-01-01'
-            }, mockContext);
-
-            const queryCall = mockAdapter.executeQuery.mock.calls[1][0] as string;
-            const queryParams = mockAdapter.executeQuery.mock.calls[1][1] as any[];
-
-            expect(queryCall).toContain('user LIKE ?');
-            expect(queryCall).toContain('event_type = ?');
-            expect(queryCall).toContain('timestamp >= ?');
-            expect(queryParams).toHaveLength(3);
-        });
-
-        it('should fallback to performance_schema if audit_log is missing', async () => {
-            // Mock table check - empty result
-            mockAdapter.executeQuery.mockResolvedValueOnce(createMockQueryResult([]));
-
-            // Mock fallback query
-            mockAdapter.executeQuery.mockResolvedValueOnce(createMockQueryResult([
-                { event: 'select' }
-            ]));
-
-            const tool = createSecurityAuditTool(mockAdapter as unknown as MySQLAdapter);
-            const result = await tool.handler({
-                limit: 10
-            }, mockContext) as { source: string };
-
-            expect(result.source).toBe('performance_schema');
-
-            const queryCall = mockAdapter.executeQuery.mock.calls[1][0] as string;
-            expect(queryCall).toContain('FROM performance_schema.events_statements_history');
-            expect(queryCall).toContain('FROM performance_schema.events_statements_history');
-        });
-
-        it('should apply user filter in fallback mode', async () => {
-            mockAdapter.executeQuery.mockResolvedValueOnce(createMockQueryResult([])); // No audit log
-            mockAdapter.executeQuery.mockResolvedValueOnce(createMockQueryResult([]));
-
-            const tool = createSecurityAuditTool(mockAdapter as unknown as MySQLAdapter);
-            await tool.handler({
-                limit: 10,
-                user: 'test_user'
-            }, mockContext);
-
-            const queryCall = mockAdapter.executeQuery.mock.calls[1][0] as string;
-            const queryParams = mockAdapter.executeQuery.mock.calls[1][1] as any[];
-
-            expect(queryCall).toContain('CURRENT_USER LIKE ?');
-            expect(queryParams[0]).toContain('test_user');
-        });
-
-        it('should handle error when checking audit log', async () => {
-            mockAdapter.executeQuery.mockRejectedValue(new Error('Access denied'));
-
-            const tool = createSecurityAuditTool(mockAdapter as unknown as MySQLAdapter);
-            const result = await tool.handler({}, mockContext) as { available: boolean };
-
-            expect(result.available).toBe(false);
-        });
+    it("should apply all filters in normal mode", async () => {
+      mockAdapter.executeQuery.mockResolvedValueOnce(
+        createMockQueryResult([{ TABLE_NAME: "audit_log" }]),
+      );
+      mockAdapter.executeQuery.mockResolvedValueOnce(createMockQueryResult([]));
+
+      const tool = createSecurityAuditTool(
+        mockAdapter as unknown as MySQLAdapter,
+      );
+      await tool.handler(
+        {
+          limit: 10,
+          user: "test",
+          eventType: "LOGIN",
+          startTime: "2023-01-01",
+        },
+        mockContext,
+      );
+
+      const queryCall = mockAdapter.executeQuery.mock.calls[1][0] as string;
+      const queryParams = mockAdapter.executeQuery.mock.calls[1][1] as any[];
+
+      expect(queryCall).toContain("user LIKE ?");
+      expect(queryCall).toContain("event_type = ?");
+      expect(queryCall).toContain("timestamp >= ?");
+      expect(queryParams).toHaveLength(3);
     });
 
-    describe('createSecurityFirewallStatusTool', () => {
-        it('should check firewall status when installed', async () => {
-            // Mock plugin check
-            mockAdapter.executeQuery.mockResolvedValueOnce(createMockQueryResult([
-                { PLUGIN_NAME: 'mysql_firewall', PLUGIN_STATUS: 'ACTIVE' }
-            ]));
-
-            // Mock variables check
-            mockAdapter.executeQuery.mockResolvedValueOnce(createMockQueryResult([
-                { Variable_name: 'mysql_firewall_mode', Value: 'ON' }
-            ]));
-
-            const tool = createSecurityFirewallStatusTool(mockAdapter as unknown as MySQLAdapter);
-            const result = await tool.handler({}, mockContext) as { installed: boolean, configuration: any };
+    it("should fallback to performance_schema if audit_log is missing", async () => {
+      // Mock table check - empty result
+      mockAdapter.executeQuery.mockResolvedValueOnce(createMockQueryResult([]));
+
+      // Mock fallback query
+      mockAdapter.executeQuery.mockResolvedValueOnce(
+        createMockQueryResult([{ event: "select" }]),
+      );
+
+      const tool = createSecurityAuditTool(
+        mockAdapter as unknown as MySQLAdapter,
+      );
+      const result = (await tool.handler(
+        {
+          limit: 10,
+        },
+        mockContext,
+      )) as { source: string };
+
+      expect(result.source).toBe("performance_schema");
+
+      const queryCall = mockAdapter.executeQuery.mock.calls[1][0] as string;
+      expect(queryCall).toContain(
+        "FROM performance_schema.events_statements_history",
+      );
+      expect(queryCall).toContain(
+        "FROM performance_schema.events_statements_history",
+      );
+    });
 
-            expect(result.installed).toBe(true);
-            expect(result.configuration).toHaveProperty('mysql_firewall_mode', 'ON');
-        });
+    it("should apply user filter in fallback mode", async () => {
+      mockAdapter.executeQuery.mockResolvedValueOnce(createMockQueryResult([])); // No audit log
+      mockAdapter.executeQuery.mockResolvedValueOnce(createMockQueryResult([]));
+
+      const tool = createSecurityAuditTool(
+        mockAdapter as unknown as MySQLAdapter,
+      );
+      await tool.handler(
+        {
+          limit: 10,
+          user: "test_user",
+        },
+        mockContext,
+      );
+
+      const queryCall = mockAdapter.executeQuery.mock.calls[1][0] as string;
+      const queryParams = mockAdapter.executeQuery.mock.calls[1][1] as any[];
+
+      expect(queryCall).toContain("CURRENT_USER LIKE ?");
+      expect(queryParams[0]).toContain("test_user");
+    });
 
-        it('should report not installed if plugin missing', async () => {
-            mockAdapter.executeQuery.mockResolvedValueOnce(createMockQueryResult([]));
+    it("should handle error when checking audit log", async () => {
+      mockAdapter.executeQuery.mockRejectedValue(new Error("Access denied"));
 
-            const tool = createSecurityFirewallStatusTool(mockAdapter as unknown as MySQLAdapter);
-            const result = await tool.handler({}, mockContext) as { installed: boolean };
+      const tool = createSecurityAuditTool(
+        mockAdapter as unknown as MySQLAdapter,
+      );
+      const result = (await tool.handler({}, mockContext)) as {
+        available: boolean;
+      };
 
-            expect(result.installed).toBe(false);
-        });
+      expect(result.available).toBe(false);
+    });
+  });
+
+  describe("createSecurityFirewallStatusTool", () => {
+    it("should check firewall status when installed", async () => {
+      // Mock plugin check
+      mockAdapter.executeQuery.mockResolvedValueOnce(
+        createMockQueryResult([
+          { PLUGIN_NAME: "mysql_firewall", PLUGIN_STATUS: "ACTIVE" },
+        ]),
+      );
+
+      // Mock variables check
+      mockAdapter.executeQuery.mockResolvedValueOnce(
+        createMockQueryResult([
+          { Variable_name: "mysql_firewall_mode", Value: "ON" },
+        ]),
+      );
+
+      const tool = createSecurityFirewallStatusTool(
+        mockAdapter as unknown as MySQLAdapter,
+      );
+      const result = (await tool.handler({}, mockContext)) as {
+        installed: boolean;
+        configuration: any;
+      };
+
+      expect(result.installed).toBe(true);
+      expect(result.configuration).toHaveProperty("mysql_firewall_mode", "ON");
     });
 
-    describe('createSecurityFirewallRulesTool', () => {
-        it('should list firewall rules', async () => {
-            // Mock users query
-            mockAdapter.executeQuery.mockResolvedValueOnce(createMockQueryResult([
-                { USERHOST: 'user@%', MODE: 'PROTECTING' }
-            ]));
-
-            // Mock rules query
-            mockAdapter.executeQuery.mockResolvedValueOnce(createMockQueryResult([
-                { USERHOST: 'user@%', RULE: 'SELECT *' }
-            ]));
+    it("should report not installed if plugin missing", async () => {
+      mockAdapter.executeQuery.mockResolvedValueOnce(createMockQueryResult([]));
 
-            const tool = createSecurityFirewallRulesTool(mockAdapter as unknown as MySQLAdapter);
-            const result = await tool.handler({
-                user: 'user'
-            }, mockContext) as { users: any[], rules: any[] };
+      const tool = createSecurityFirewallStatusTool(
+        mockAdapter as unknown as MySQLAdapter,
+      );
+      const result = (await tool.handler({}, mockContext)) as {
+        installed: boolean;
+      };
 
-            expect(mockAdapter.executeQuery).toHaveBeenCalledTimes(2);
-            expect(result.users).toHaveLength(1);
-            expect(result.rules).toHaveLength(1);
-        });
+      expect(result.installed).toBe(false);
+    });
+  });
+
+  describe("createSecurityFirewallRulesTool", () => {
+    it("should list firewall rules", async () => {
+      // Mock users query
+      mockAdapter.executeQuery.mockResolvedValueOnce(
+        createMockQueryResult([{ USERHOST: "user@%", MODE: "PROTECTING" }]),
+      );
+
+      // Mock rules query
+      mockAdapter.executeQuery.mockResolvedValueOnce(
+        createMockQueryResult([{ USERHOST: "user@%", RULE: "SELECT *" }]),
+      );
+
+      const tool = createSecurityFirewallRulesTool(
+        mockAdapter as unknown as MySQLAdapter,
+      );
+      const result = (await tool.handler(
+        {
+          user: "user",
+        },
+        mockContext,
+      )) as { users: any[]; rules: any[] };
+
+      expect(mockAdapter.executeQuery).toHaveBeenCalledTimes(2);
+      expect(result.users).toHaveLength(1);
+      expect(result.rules).toHaveLength(1);
+    });
 
-        it('should handle errors gracefully', async () => {
-            mockAdapter.executeQuery.mockRejectedValue(new Error('Table missing'));
+    it("should handle errors gracefully", async () => {
+      mockAdapter.executeQuery.mockRejectedValue(new Error("Table missing"));
 
-            const tool = createSecurityFirewallRulesTool(mockAdapter as unknown as MySQLAdapter);
-            const result = await tool.handler({}, mockContext) as { available: boolean };
+      const tool = createSecurityFirewallRulesTool(
+        mockAdapter as unknown as MySQLAdapter,
+      );
+      const result = (await tool.handler({}, mockContext)) as {
+        available: boolean;
+      };
 
-            expect(result.available).toBe(false);
-        });
+      expect(result.available).toBe(false);
     });
+  });
 });

package/src/adapters/mysql/tools/security/__tests__/data-protection.test.ts

@@ -1,111 +1,167 @@
 /**
  * mysql-mcp - Security Data Protection Tools Unit Tests
- * 
+ *
  * Comprehensive tests for data-protection.ts module functions.
  */
 
-import { describe, it, expect, vi, beforeEach } from 'vitest';
+import { describe, it, expect, vi, beforeEach } from "vitest";
 import {
-    createSecurityMaskDataTool,
-    createSecurityUserPrivilegesTool,
-    createSecuritySensitiveTablesTool
-} from '../data-protection.js';
-import type { MySQLAdapter } from '../../MySQLAdapter.js';
-import { createMockMySQLAdapter, createMockRequestContext, createMockQueryResult } from '../../../../../__tests__/mocks/index.js';
-
-describe('Security Data Protection Tools', () => {
-    let mockAdapter: ReturnType<typeof createMockMySQLAdapter>;
-    let mockContext: ReturnType<typeof createMockRequestContext>;
-
-    beforeEach(() => {
-        vi.clearAllMocks();
-        mockAdapter = createMockMySQLAdapter();
-        mockContext = createMockRequestContext();
+  createSecurityMaskDataTool,
+  createSecurityUserPrivilegesTool,
+  createSecuritySensitiveTablesTool,
+} from "../data-protection.js";
+import type { MySQLAdapter } from "../../../MySQLAdapter.js";
+import {
+  createMockMySQLAdapter,
+  createMockRequestContext,
+  createMockQueryResult,
+} from "../../../../../__tests__/mocks/index.js";
+
+describe("Security Data Protection Tools", () => {
+  let mockAdapter: ReturnType<typeof createMockMySQLAdapter>;
+  let mockContext: ReturnType<typeof createMockRequestContext>;
+
+  beforeEach(() => {
+    vi.clearAllMocks();
+    mockAdapter = createMockMySQLAdapter();
+    mockContext = createMockRequestContext();
+  });
+
+  describe("createSecurityMaskDataTool", () => {
+    it("should mask email addresses", async () => {
+      const tool = createSecurityMaskDataTool(
+        mockAdapter as unknown as MySQLAdapter,
+      );
+      const result = (await tool.handler(
+        {
+          value: "john.doe@example.com",
+          type: "email",
+        },
+        mockContext,
+      )) as { masked: string };
+
+      expect(result.masked).toBe("j******e@example.com");
     });
 
-    describe('createSecurityMaskDataTool', () => {
-        it('should mask email addresses', async () => {
-            const tool = createSecurityMaskDataTool(mockAdapter as unknown as MySQLAdapter);
-            const result = await tool.handler({
-                value: 'john.doe@example.com',
-                type: 'email'
-            }, mockContext) as { masked: string };
-
-            expect(result.masked).toBe('j******e@example.com');
-        });
-
-        it('should mask phone numbers', async () => {
-            const tool = createSecurityMaskDataTool(mockAdapter as unknown as MySQLAdapter);
-            const result = await tool.handler({
-                value: '555-123-4567',
-                type: 'phone'
-            }, mockContext) as { masked: string };
-
-            expect(result.masked).toBe('******4567');
-        });
-
-        it('should mask partial text', async () => {
-            const tool = createSecurityMaskDataTool(mockAdapter as unknown as MySQLAdapter);
-            const result = await tool.handler({
-                value: 'sensitive data',
-                type: 'partial',
-                keepFirst: 2,
-                keepLast: 2
-            }, mockContext) as { masked: string };
-
-            expect(result.masked).toBe('se**********ta');
-        });
+    it("should mask phone numbers", async () => {
+      const tool = createSecurityMaskDataTool(
+        mockAdapter as unknown as MySQLAdapter,
+      );
+      const result = (await tool.handler(
+        {
+          value: "555-123-4567",
+          type: "phone",
+        },
+        mockContext,
+      )) as { masked: string };
+
+      expect(result.masked).toBe("******4567");
     });
 
-    describe('createSecurityUserPrivilegesTool', () => {
-        it('should get user privileges and roles', async () => {
-            // Mock users query
-            mockAdapter.executeQuery.mockResolvedValueOnce(createMockQueryResult([
-                { User: 'john', Host: 'localhost', authPlugin: 'native', accountLocked: 'N', passwordExpired: 'N' }
-            ]));
-
-            // Mock permissions query for john
-            mockAdapter.executeQuery.mockResolvedValueOnce(createMockQueryResult([
-                { 'Grants for john@localhost': 'GRANT ALL PRIVILEGES ON *.* TO \'john\'@\'localhost\'' }
-            ]));
-
-            // Mock roles query
-            mockAdapter.executeQuery.mockResolvedValueOnce(createMockQueryResult([
-                { FROM_USER: 'admin_role', FROM_HOST: '%' }
-            ]));
-
-
-            const tool = createSecurityUserPrivilegesTool(mockAdapter as unknown as MySQLAdapter);
-            const result = await tool.handler({
-                user: 'john'
-            }, mockContext) as { users: any[] };
-
-            expect(mockAdapter.executeQuery).toHaveBeenCalledTimes(3);
-            expect(result.users).toHaveLength(1);
-            expect(result.users[0].user).toBe('john');
-            expect(result.users[0].grants).toHaveLength(1);
-            expect(result.users[0].roles).toContain('admin_role@%');
-        });
+    it("should mask partial text", async () => {
+      const tool = createSecurityMaskDataTool(
+        mockAdapter as unknown as MySQLAdapter,
+      );
+      const result = (await tool.handler(
+        {
+          value: "sensitive data",
+          type: "partial",
+          keepFirst: 2,
+          keepLast: 2,
+        },
+        mockContext,
+      )) as { masked: string };
+
+      expect(result.masked).toBe("se**********ta");
     });
-
-    describe('createSecuritySensitiveTablesTool', () => {
-        it('should identify sensitive tables', async () => {
-            mockAdapter.executeQuery.mockResolvedValue(createMockQueryResult([
-                { tableName: 'users', columnName: 'password_hash', dataType: 'varchar' },
-                { tableName: 'users', columnName: 'email', dataType: 'varchar' },
-                { tableName: 'payments', columnName: 'credit_card', dataType: 'varchar' }
-            ]));
-
-            const tool = createSecuritySensitiveTablesTool(mockAdapter as unknown as MySQLAdapter);
-            const result = await tool.handler({
-                schema: 'test_db'
-            }, mockContext) as { sensitiveTables: any[], totalSensitiveColumns: number };
-
-            expect(mockAdapter.executeQuery).toHaveBeenCalled();
-            expect(result.sensitiveTables).toHaveLength(2); // users, payments
-            expect(result.totalSensitiveColumns).toBe(3);
-            expect(result.sensitiveTables[0].table).toBe('users');
-            expect(result.sensitiveTables[0].sensitiveColumns).toHaveLength(2);
-        });
+  });
+
+  describe("createSecurityUserPrivilegesTool", () => {
+    it("should get user privileges and roles", async () => {
+      // Mock P154 user existence pre-check
+      mockAdapter.executeQuery.mockResolvedValueOnce(
+        createMockQueryResult([{ User: "john", Host: "localhost" }]),
+      );
+
+      // Mock users query
+      mockAdapter.executeQuery.mockResolvedValueOnce(
+        createMockQueryResult([
+          {
+            User: "john",
+            Host: "localhost",
+            authPlugin: "native",
+            accountLocked: "N",
+            passwordExpired: "N",
+          },
+        ]),
+      );
+
+      // Mock permissions query for john
+      mockAdapter.executeQuery.mockResolvedValueOnce(
+        createMockQueryResult([
+          {
+            "Grants for john@localhost":
+              "GRANT ALL PRIVILEGES ON *.* TO 'john'@'localhost'",
+          },
+        ]),
+      );
+
+      // Mock roles query
+      mockAdapter.executeQuery.mockResolvedValueOnce(
+        createMockQueryResult([{ FROM_USER: "admin_role", FROM_HOST: "%" }]),
+      );
+
+      const tool = createSecurityUserPrivilegesTool(
+        mockAdapter as unknown as MySQLAdapter,
+      );
+      const result = (await tool.handler(
+        {
+          user: "john",
+        },
+        mockContext,
+      )) as { users: any[] };
+
+      expect(mockAdapter.executeQuery).toHaveBeenCalledTimes(4);
+      expect(result.users).toHaveLength(1);
+      expect(result.users[0].user).toBe("john");
+      expect(result.users[0].grants).toHaveLength(1);
+      expect(result.users[0].roles).toContain("admin_role@%");
+    });
+  });
+
+  describe("createSecuritySensitiveTablesTool", () => {
+    it("should identify sensitive tables", async () => {
+      mockAdapter.executeQuery.mockResolvedValue(
+        createMockQueryResult([
+          {
+            tableName: "users",
+            columnName: "password_hash",
+            dataType: "varchar",
+          },
+          { tableName: "users", columnName: "email", dataType: "varchar" },
+          {
+            tableName: "payments",
+            columnName: "credit_card",
+            dataType: "varchar",
+          },
+        ]),
+      );
+
+      const tool = createSecuritySensitiveTablesTool(
+        mockAdapter as unknown as MySQLAdapter,
+      );
+      const result = (await tool.handler(
+        {
+          schema: "test_db",
+        },
+        mockContext,
+      )) as { sensitiveTables: any[]; totalSensitiveColumns: number };
+
+      expect(mockAdapter.executeQuery).toHaveBeenCalled();
+      expect(result.sensitiveTables).toHaveLength(2); // users, payments
+      expect(result.totalSensitiveColumns).toBe(3);
+      expect(result.sensitiveTables[0].table).toBe("users");
+      expect(result.sensitiveTables[0].sensitiveColumns).toHaveLength(2);
     });
+  });
 });