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 => ({ 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 let targets: DriverTarget[] let tunnels: FakeTunnel[] let createDriver: (target: DriverTarget) => Driver let createTunnel: ReturnType 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((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((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('retries instead of throwing when a build fails after ownership is lost', async () => { let release: () => void = () => {} const gate = new Promise((resolve) => { release = resolve }) // First build is gated and then rejects; later builds succeed. const gatedTunnel = vi .fn(createTunnel) .mockImplementationOnce(async () => { await gate throw new Error('stale build boom') }) .mockImplementation(createTunnel) resolved = { config: config(sshExtra()), source: 'store', hash: 'ssh1' } const manager = createManager(registry, { createDriver, createTunnel: gatedTunnel }) const pending = manager.get('c') // Drop ownership of the in-flight slot (synchronous cache.delete) before // its build rejects; do not await invalidate or it blocks on the gate. const invalidated = manager.invalidate('c') release() // The stale build's error is swallowed; get retries and resolves fresh. const managed = await pending expect(managed.driver).toBe(drivers[0]) await invalidated }) 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]) }) it('disposeAll waits for an in-flight rotation dispose', async () => { let release: () => void = () => {} const gate = new Promise((resolve) => { release = resolve }) // The first driver's dispose is gated; rotation fires it and forgets it. const gatingCreateDriver = (target: DriverTarget): Driver => { targets.push(target) const first = drivers.length === 0 const driver = fakeDriver() if (first) { driver.dispose = vi.fn(async () => { await gate driver.disposed = true }) } drivers.push(driver) return driver } const manager = createManager(registry, { createDriver: gatingCreateDriver, createTunnel }) await manager.get('c') // Hash change rotates the old entry out: its dispose is fire-and-forget. resolved = { config: config({ port: 9999 }), source: 'store', hash: 'h2' } await manager.get('c') expect(drivers[0].disposed).toBe(false) let resolvedAll = false const all = manager.disposeAll().then(() => { resolvedAll = true }) // Flush all currently-queued microtasks: disposeAll must still be // pending because the gated rotation dispose has not completed. for (let i = 0; i < 10; i += 1) { await Promise.resolve() } expect(resolvedAll).toBe(false) expect(drivers[0].disposed).toBe(false) release() await all // disposeAll only resolves once the rotated-out driver finished disposing. expect(drivers[0].disposed).toBe(true) expect(drivers[1].disposed).toBe(true) }) })