dbmole-mcp/test/unit/tools/connections.test.ts

import { mkdtempSync, rmSync } from 'node:fs'
import { tmpdir } from 'node:os'
import { join } from 'node:path'
import { McpServer } from '@modelcontextprotocol/sdk/server/mcp.js'
import { afterEach, beforeEach, describe, expect, it, vi } from 'vitest'
import { createRegistry } from '../../../src/config/registry.js'
import { writeStore } from '../../../src/config/store.js'
import type { ConnectionConfig } from '../../../src/config/types.js'
import { registerConnectionTools } from '../../../src/tools/connections.js'
import { callTool, connectClient, type FakeManager, fakeManager } from '../helpers.js'

const conn = (name: string, extra: Partial<ConnectionConfig> = {}): ConnectionConfig => ({
    name,
    type: 'postgres',
    host: 'localhost',
    user: 'postgres',
    readonly: false,
    ...extra
})

describe('connection tools', () => {
    let dir: string
    let storePath: string
    let manager: FakeManager
    let client: Awaited<ReturnType<typeof connectClient>>

    beforeEach(async () => {
        dir = mkdtempSync(join(tmpdir(), 'dbmole-tools-'))
        storePath = join(dir, 'connections.json')
        writeStore(storePath, [conn('existing', { password: 'secret', database: 'app' })])
        const registry = createRegistry({ storePath, env: {} })
        manager = fakeManager()
        const server = new McpServer({ name: 'test', version: '0.0.0' })
        registerConnectionTools(server, registry, manager)
        client = await connectClient(server)
    })

    afterEach(() => {
        rmSync(dir, { recursive: true, force: true })
    })

    it('list_connections returns public view without secrets', async () => {
        const response = await callTool(client, 'list_connections')
        expect(response.isError).toBe(false)
        const list = response.json() as Array<Record<string, unknown>>
        expect(list).toEqual([
            {
                name: 'existing',
                type: 'postgres',
                host: 'localhost',
                port: 5432,
                database: 'app',
                readonly: false,
                source: 'store',
                ssh: false
            }
        ])
        expect(response.text).not.toContain('secret')
    })

    it('add_connection persists and reports the new connection', async () => {
        const response = await callTool(client, 'add_connection', {
            name: 'fresh',
            type: 'mysql',
            host: 'db',
            user: 'root'
        })
        expect(response.isError).toBe(false)
        const list = await callTool(client, 'list_connections')
        expect((list.json() as Array<{ name: string }>).map((c) => c.name)).toEqual([
            'existing',
            'fresh'
        ])
    })

    it('add_connection rejects unknown fields', async () => {
        const response = await callTool(client, 'add_connection', {
            name: 'typo',
            type: 'postgres',
            host: 'h',
            user: 'u',
            readOnly: true
        })
        expect(response.isError).toBe(true)
        const list = await callTool(client, 'list_connections')
        const names = (list.json() as Array<{ name: string }>).map((c) => c.name)
        expect(names).not.toContain('typo')
    })

    it('add_connection rejects duplicates with isError', async () => {
        const response = await callTool(client, 'add_connection', {
            name: 'existing',
            type: 'postgres',
            host: 'x',
            user: 'u'
        })
        expect(response.isError).toBe(true)
        expect(response.text).toContain('already exists')
    })

    it('update_connection patches and invalidates the manager cache', async () => {
        const response = await callTool(client, 'update_connection', {
            name: 'existing',
            patch: { host: 'new-host', database: null }
        })
        expect(response.isError).toBe(false)
        expect(manager.invalidate).toHaveBeenCalledWith('existing')
        const list = await callTool(client, 'list_connections')
        const updated = (list.json() as Array<Record<string, unknown>>)[0]
        expect(updated.host).toBe('new-host')
        expect(updated.database).toBeNull()
    })

    it('update_connection surfaces validation errors', async () => {
        const response = await callTool(client, 'update_connection', {
            name: 'existing',
            patch: { type: 'oracle' }
        })
        expect(response.isError).toBe(true)
    })

    it('remove_connection deletes and invalidates', async () => {
        const response = await callTool(client, 'remove_connection', { name: 'existing' })
        expect(response.isError).toBe(false)
        expect(manager.invalidate).toHaveBeenCalledWith('existing')
        const list = await callTool(client, 'list_connections')
        expect(list.json()).toEqual([])
    })

    it('remove_connection of unknown name reports available connections', async () => {
        const response = await callTool(client, 'remove_connection', { name: 'ghost' })
        expect(response.isError).toBe(true)
        expect(response.text).toContain('existing')
    })

    it('test_connection reports version and latency', async () => {
        manager.get.mockResolvedValue({
            driver: { serverVersion: vi.fn(async () => '17.2') },
            config: conn('existing'),
            source: 'store'
        })
        const response = await callTool(client, 'test_connection', { name: 'existing' })
        expect(response.isError).toBe(false)
        const payload = response.json() as Record<string, unknown>
        expect(payload.ok).toBe(true)
        expect(payload.version).toBe('17.2')
        expect(typeof payload.latencyMs).toBe('number')
    })

    it('test_connection formats driver failures', async () => {
        manager.get.mockResolvedValue({
            driver: {
                serverVersion: vi.fn(async () => {
                    throw Object.assign(new Error('connect ECONNREFUSED'), { code: 'ECONNREFUSED' })
                })
            },
            config: conn('existing'),
            source: 'store'
        })
        const response = await callTool(client, 'test_connection', { name: 'existing' })
        expect(response.isError).toBe(true)
        expect(response.text).toContain('ECONNREFUSED')
    })

    it('test_connection reports manager failures (e.g. tunnel)', async () => {
        manager.get.mockRejectedValue(new Error('ssh auth failed'))
        const response = await callTool(client, 'test_connection', { name: 'existing' })
        expect(response.isError).toBe(true)
        expect(response.text).toContain('ssh auth failed')
    })
})