dbmole-mcp/packages/esmole-mcp/test/unit/tools/register.test.ts

import type { McpServer } from '@modelcontextprotocol/sdk/server/mcp.js'
import type { CallToolResult } from '@modelcontextprotocol/sdk/types.js'
import { describe, expect, it } from 'vitest'
import type { EsBackend } from '../../../src/es/backend.js'
import type { EsManaged, EsManager } from '../../../src/manager.js'
import { registerEsTools } from '../../../src/tools/index.js'

type RequestCall = {
    method: string
    path: string
    opts?: { body?: unknown; params?: Record<string, string | number | boolean> }
}

type ToolHandler = (args: Record<string, unknown>) => Promise<CallToolResult>

type RecordingServer = {
    handlers: Map<string, ToolHandler>
    server: McpServer
}

// Only registerTool is exercised; expose just that method and cast once to
// McpServer (AGENTS.md allows a single cast for the server/manager shape —
// no `as unknown as`).
const fakeServer = (): RecordingServer => {
    const handlers = new Map<string, ToolHandler>()
    const recorder = {
        registerTool: (name: string, _config: unknown, handler: ToolHandler): void => {
            handlers.set(name, handler)
        }
    }
    return { handlers, server: recorder as McpServer }
}

const fakeManager = (calls: RequestCall[], body: unknown): EsManager => {
    const backend: EsBackend = {
        request: async (method, path, opts) => {
            calls.push({ method, path, opts })
            return { status: 200, body }
        },
        dispose: async () => {}
    }
    const managed: EsManaged = {
        backend,
        config: { name: 'c', readonly: false } as EsManaged['config'],
        source: 'test'
    }
    return {
        get: async () => managed,
        invalidate: async () => {},
        disposeAll: async () => {}
    }
}

const isCallToolResult = (res: CallToolResult): boolean => Array.isArray(res.content)

describe('registerEsTools', () => {
    it('registers all five tools', () => {
        const { handlers, server } = fakeServer()
        registerEsTools(server, fakeManager([], {}))
        expect([...handlers.keys()].sort()).toEqual(
            [
                'es_cluster_health',
                'es_get_mapping',
                'es_list_indices',
                'es_request',
                'es_search'
            ].sort()
        )
    })

    it('es_request handler calls backend with method/path and returns a CallToolResult', async () => {
        const calls: RequestCall[] = []
        const { handlers, server } = fakeServer()
        registerEsTools(server, fakeManager(calls, { ok: true }))
        const handler = handlers.get('es_request')
        if (handler === undefined) {
            throw new Error('es_request not registered')
        }
        const res = await handler({
            connection: 'c',
            method: 'POST',
            path: '/idx/_doc',
            body: { a: 1 }
        })
        expect(calls[0].method).toBe('POST')
        expect(calls[0].path).toBe('/idx/_doc')
        expect(isCallToolResult(res)).toBe(true)
        expect(res.isError).toBeFalsy()
    })

    it('es_search handler posts /{index}/_search', async () => {
        const calls: RequestCall[] = []
        const { handlers, server } = fakeServer()
        registerEsTools(server, fakeManager(calls, { hits: {} }))
        const handler = handlers.get('es_search')
        if (handler === undefined) {
            throw new Error('es_search not registered')
        }
        const res = await handler({ connection: 'c', index: 'books', query: { match_all: {} } })
        expect(calls[0].method).toBe('POST')
        expect(calls[0].path).toBe('/books/_search')
        expect(isCallToolResult(res)).toBe(true)
    })

    it('es_list_indices handler GETs /_cat/indices/<pattern> with format json', async () => {
        const calls: RequestCall[] = []
        const { handlers, server } = fakeServer()
        registerEsTools(server, fakeManager(calls, []))
        const handler = handlers.get('es_list_indices')
        if (handler === undefined) {
            throw new Error('es_list_indices not registered')
        }
        const res = await handler({ connection: 'c', pattern: 'logs-*' })
        expect(calls[0].method).toBe('GET')
        expect(calls[0].path).toBe('/_cat/indices/logs-*')
        expect(calls[0].opts?.params).toEqual({ format: 'json' })
        expect(isCallToolResult(res)).toBe(true)
    })

    it('es_get_mapping handler GETs /{index}/_mapping and flattens by default', async () => {
        const calls: RequestCall[] = []
        const { handlers, server } = fakeServer()
        registerEsTools(
            server,
            fakeManager(calls, { idx: { mappings: { properties: { y: { type: 'long' } } } } })
        )
        const handler = handlers.get('es_get_mapping')
        if (handler === undefined) {
            throw new Error('es_get_mapping not registered')
        }
        const res = await handler({ connection: 'c', index: 'idx', raw: false })
        expect(calls[0].method).toBe('GET')
        expect(calls[0].path).toBe('/idx/_mapping')
        expect(isCallToolResult(res)).toBe(true)
        const payload = JSON.parse((res.content as Array<{ text: string }>)[0].text)
        expect(payload.fields).toContainEqual({ field: 'y', type: 'long' })
    })

    it('es_cluster_health handler GETs /_cluster/health', async () => {
        const calls: RequestCall[] = []
        const { handlers, server } = fakeServer()
        registerEsTools(server, fakeManager(calls, { status: 'green' }))
        const handler = handlers.get('es_cluster_health')
        if (handler === undefined) {
            throw new Error('es_cluster_health not registered')
        }
        const res = await handler({ connection: 'c' })
        expect(calls[0].method).toBe('GET')
        expect(calls[0].path).toBe('/_cluster/health')
        expect(isCallToolResult(res)).toBe(true)
    })
})