dbmole-mcp/test/unit/db/manager.test.ts

import { beforeEach, describe, expect, it, vi } from 'vitest'
import type { Registry } from '../../../src/config/registry.js'
import type { ConnectionConfig, ResolvedConnection, SshConfig } from '../../../src/config/types.js'
import type { Driver, DriverTarget } from '../../../src/db/driver.js'
import { createManager } from '../../../src/db/manager.js'
import type { Tunnel } from '../../../src/net/tunnel.js'

const config = (extra: Partial<ConnectionConfig> = {}): ConnectionConfig => ({
    name: 'c',
    type: 'postgres',
    host: 'db-host',
    user: 'u',
    readonly: false,
    ...extra
})

const fakeDriver = (): Driver & { disposed: boolean } => {
    const driver = {
        disposed: false,
        query: vi.fn(),
        listDatabases: vi.fn(),
        listTables: vi.fn(),
        describeTable: vi.fn(),
        serverVersion: vi.fn(),
        dispose: vi.fn(async () => {
            driver.disposed = true
        })
    }
    return driver
}

type FakeTunnel = Tunnel & { closed: boolean }

describe('createManager', () => {
    let resolved: ResolvedConnection
    let registry: Registry
    let drivers: Array<Driver & { disposed: boolean }>
    let targets: DriverTarget[]
    let tunnels: FakeTunnel[]
    let createDriver: (target: DriverTarget) => Driver
    let createTunnel: ReturnType<typeof vi.fn>

    beforeEach(() => {
        resolved = { config: config(), source: 'store', hash: 'h1' }
        registry = {
            list: vi.fn(() => [resolved]),
            resolve: vi.fn(() => resolved),
            add: vi.fn(),
            update: vi.fn(),
            remove: vi.fn()
        }
        drivers = []
        targets = []
        tunnels = []
        createDriver = (target: DriverTarget) => {
            targets.push(target)
            const driver = fakeDriver()
            drivers.push(driver)
            return driver
        }
        createTunnel = vi.fn(async () => {
            const tunnel: FakeTunnel = {
                localHost: '127.0.0.1',
                localPort: 54321,
                closed: false,
                isClosed: () => tunnel.closed,
                close: async () => {
                    tunnel.closed = true
                }
            }
            tunnels.push(tunnel)
            return tunnel
        })
    })

    const sshExtra = (): { ssh: SshConfig } => ({
        ssh: { host: 'bastion', port: 22, user: 'root', password: 'x' }
    })

    it('builds lazily and caches by hash', async () => {
        const manager = createManager(registry, { createDriver, createTunnel })
        const first = await manager.get('c')
        const second = await manager.get('c')
        expect(first.driver).toBe(second.driver)
        expect(drivers).toHaveLength(1)
        expect(targets[0]).toMatchObject({ host: 'db-host', port: 5432 })
        expect(createTunnel).not.toHaveBeenCalled()
    })

    it('rebuilds and disposes the old driver when the hash changes', async () => {
        const manager = createManager(registry, { createDriver, createTunnel })
        const first = await manager.get('c')
        resolved = { config: config({ port: 9999 }), source: 'store', hash: 'h2' }
        const second = await manager.get('c')
        expect(second.driver).not.toBe(first.driver)
        expect(drivers[0].disposed).toBe(true)
        expect(targets[1]).toMatchObject({ port: 9999 })
    })

    it('routes through the tunnel endpoint when ssh is configured', async () => {
        resolved = { config: config(sshExtra()), source: 'store', hash: 'ssh1' }
        const manager = createManager(registry, { createDriver, createTunnel })
        await manager.get('c')
        expect(createTunnel).toHaveBeenCalledWith(
            expect.objectContaining({ host: 'bastion' }),
            'db-host',
            5432
        )
        expect(targets[0]).toMatchObject({ host: '127.0.0.1', port: 54321 })
    })

    it('uses the engine default port for mysql', async () => {
        resolved = { config: config({ type: 'mysql' }), source: 'store', hash: 'm1' }
        const manager = createManager(registry, { createDriver, createTunnel })
        await manager.get('c')
        expect(targets[0]).toMatchObject({ port: 3306 })
    })

    it('rebuilds when the tunnel reports closed', async () => {
        resolved = { config: config(sshExtra()), source: 'store', hash: 'ssh1' }
        const manager = createManager(registry, { createDriver, createTunnel })
        const first = await manager.get('c')
        tunnels[0].closed = true
        const second = await manager.get('c')
        expect(second.driver).not.toBe(first.driver)
        expect(drivers[0].disposed).toBe(true)
        expect(tunnels).toHaveLength(2)
        expect(second.driver).toBe(drivers[1])
    })

    it('deduplicates concurrent builds', async () => {
        const manager = createManager(registry, { createDriver, createTunnel })
        const [a, b] = await Promise.all([manager.get('c'), manager.get('c')])
        expect(a.driver).toBe(b.driver)
        expect(drivers).toHaveLength(1)
    })

    it('clears the cache when a build fails, allowing retry', async () => {
        const failing = vi
            .fn()
            .mockRejectedValueOnce(new Error('tunnel down'))
            .mockImplementation(createTunnel)
        resolved = { config: config(sshExtra()), source: 'store', hash: 'ssh1' }
        const manager = createManager(registry, { createDriver, createTunnel: failing })
        await expect(manager.get('c')).rejects.toThrow('tunnel down')
        const retried = await manager.get('c')
        expect(retried.driver).toBe(drivers[0])
    })

    it('invalidate disposes driver and tunnel', async () => {
        resolved = { config: config(sshExtra()), source: 'store', hash: 'ssh1' }
        const manager = createManager(registry, { createDriver, createTunnel })
        await manager.get('c')
        await manager.invalidate('c')
        expect(drivers[0].disposed).toBe(true)
        expect(tunnels[0].closed).toBe(true)
    })

    it('invalidate of an unknown name is a no-op', async () => {
        const manager = createManager(registry, { createDriver, createTunnel })
        await expect(manager.invalidate('ghost')).resolves.toBeUndefined()
    })

    it('invalidate during an in-flight build disposes the entry once built', async () => {
        let release: () => void = () => {}
        const gate = new Promise<void>((resolve) => {
            release = resolve
        })
        const gatedTunnel = vi.fn(async (ssh: SshConfig, host: string, port: number) => {
            await gate
            return createTunnel(ssh, host, port)
        })
        resolved = { config: config(sshExtra()), source: 'store', hash: 'ssh1' }
        const manager = createManager(registry, { createDriver, createTunnel: gatedTunnel })
        const pending = manager.get('c')
        const invalidated = manager.invalidate('c')
        release()
        await invalidated
        // The first build is disposed by invalidate; the pending get loses
        // ownership and retries, resolving to a freshly rebuilt driver.
        expect((await pending).driver).toBe(drivers[1])
        expect(drivers[0].disposed).toBe(true)
        expect(tunnels[0].closed).toBe(true)
    })

    it('get racing disposeAll does not return a disposed entry', async () => {
        let release: () => void = () => {}
        const gate = new Promise<void>((resolve) => {
            release = resolve
        })
        const gatedTunnel = vi.fn(async (ssh: SshConfig, host: string, port: number) => {
            await gate
            return createTunnel(ssh, host, port)
        })
        resolved = { config: config(sshExtra()), source: 'store', hash: 'ssh1' }
        const manager = createManager(registry, { createDriver, createTunnel: gatedTunnel })
        const pending = manager.get('c')
        const disposed = manager.disposeAll()
        release()
        await disposed
        // The gated build is disposed; the pending get retries and returns a
        // fresh entry that was never disposed.
        const managed = await pending
        expect(managed.driver).toBe(drivers[1])
        expect(drivers[0].disposed).toBe(true)
        expect(drivers[1].disposed).toBe(false)
    })

    it('closes the tunnel when createDriver throws and a retry succeeds', async () => {
        const throwingDriver = vi
            .fn<(target: DriverTarget) => Driver>()
            .mockImplementationOnce(() => {
                throw new Error('driver boom')
            })
            .mockImplementation(createDriver)
        resolved = { config: config(sshExtra()), source: 'store', hash: 'ssh1' }
        const manager = createManager(registry, {
            createDriver: throwingDriver,
            createTunnel
        })
        await expect(manager.get('c')).rejects.toThrow('driver boom')
        expect(tunnels[0].closed).toBe(true)
        const retried = await manager.get('c')
        expect(retried.driver).toBe(drivers[0])
        expect(tunnels[1].closed).toBe(false)
    })

    it('disposeAll disposes everything', async () => {
        const manager = createManager(registry, { createDriver, createTunnel })
        await manager.get('c')
        await manager.disposeAll()
        expect(drivers[0].disposed).toBe(true)
        const again = await manager.get('c')
        expect(again.driver).toBe(drivers[1])
    })
})