sh-mcp/packages/knowledge-base/test/knowledge-base.test.ts

/**
 * Unit tests for @sh-mcp/knowledge-base.
 *
 * All tests use:
 *   - A mock AuthContext with the required ops:read scope.
 *   - A mock KnowledgeBaseClient — no AWS calls, no network.
 *
 * Coverage targets:
 *   - Happy path: results returned and shaped correctly.
 *   - Empty result: KB returns no matches.
 *   - Scope enforcement: missing scope throws ScopeError.
 *   - Client error: upstream error surfaces as a rejected promise.
 *   - Throttle / retry: upstream ThrottlingException propagates (retry
 *     logic, if added, would be tested here).
 */

import { describe, it, expect, vi, beforeEach } from 'vitest';
import { createKnowledgeBaseTools } from '../src/tools.js';
import type { KnowledgeBaseClient, KnowledgeBaseResult } from '../src/client.js';
import type { AuthContext } from '@sh-mcp/shared';

// ---------------------------------------------------------------------------
// Test fixtures
// ---------------------------------------------------------------------------

/** A valid AuthContext carrying the ops:read scope. */
const authorisedCtx: AuthContext = {
  sub: 'lauren@seahavenind.com',
  scopes: ['ops:read'],
  aud: 'sh-mcp-ops',
};

/** An AuthContext with no scopes — used to test scope enforcement. */
const unauthorisedCtx: AuthContext = {
  sub: 'guest@example.com',
  scopes: [],
  aud: 'sh-mcp-ops',
};

/** A finance-only AuthContext (finance:read but not ops:read). */
const financeOnlyCtx: AuthContext = {
  sub: 'accounting@seahavenind.com',
  scopes: ['finance:read'],
  aud: 'sh-mcp-finance',
};

/** Sample KB results returned by the mock. */
const sampleResults: KnowledgeBaseResult[] = [
  {
    source: 's3://sh-kb-data/notion/procedures.md',
    score: 0.92,
    passage: 'All maintenance requests must be submitted via the work-order portal.',
  },
  {
    source: 's3://sh-kb-data/work-orders/WO-1042.md',
    score: 0.78,
    passage: 'Work order 1042: HVAC inspection completed 2025-11-15.',
  },
];

// ---------------------------------------------------------------------------
// Mock client factory
// ---------------------------------------------------------------------------

function makeMockClient(
  implementation?: Partial<KnowledgeBaseClient>
): KnowledgeBaseClient {
  return {
    retrieve: vi.fn().mockResolvedValue(sampleResults),
    ...implementation,
  };
}

// ---------------------------------------------------------------------------
// Tests
// ---------------------------------------------------------------------------

describe('search_knowledge_base', () => {
  let mockClient: KnowledgeBaseClient;

  beforeEach(() => {
    mockClient = makeMockClient();
  });

  // -------------------------------------------------------------------------
  // Happy path
  // -------------------------------------------------------------------------

  it('returns shaped results for a valid query', async () => {
    const [tool] = createKnowledgeBaseTools(mockClient);

    const output = await tool.handler(
      { query: 'maintenance procedures', maxResults: 5 },
      authorisedCtx
    );

    expect(output.count).toBe(2);
    expect(output.results).toHaveLength(2);

    const [first] = output.results;
    expect(first.source).toBe('s3://sh-kb-data/notion/procedures.md');
    expect(first.score).toBe(0.92);
    expect(first.passage).toContain('maintenance requests');
  });

  it('passes query and maxResults through to the client', async () => {
    const [tool] = createKnowledgeBaseTools(mockClient);

    await tool.handler({ query: 'HVAC vendors', maxResults: 3 }, authorisedCtx);

    expect(mockClient.retrieve).toHaveBeenCalledOnce();
    expect(mockClient.retrieve).toHaveBeenCalledWith({
      query: 'HVAC vendors',
      maxResults: 3,
    });
  });

  it('defaults maxResults to 5 when omitted', async () => {
    const [tool] = createKnowledgeBaseTools(mockClient);

    await tool.handler({ query: 'fire safety' }, authorisedCtx);

    expect(mockClient.retrieve).toHaveBeenCalledWith({
      query: 'fire safety',
      maxResults: 5,
    });
  });

  // -------------------------------------------------------------------------
  // Tool metadata assertions
  // -------------------------------------------------------------------------

  it('has correct tool metadata', () => {
    const [tool] = createKnowledgeBaseTools(mockClient);

    expect(tool.name).toBe('search_knowledge_base');
    expect(tool.tier).toBe('ops');
    expect(tool.requiredScope).toBe('ops:read');
    expect(tool.description).toMatch(/knowledge base/i);
  });

  it('exports exactly one tool', () => {
    const tools = createKnowledgeBaseTools(mockClient);
    expect(tools).toHaveLength(1);
  });

  // -------------------------------------------------------------------------
  // Empty result
  // -------------------------------------------------------------------------

  it('returns an empty results array when the KB finds no matches', async () => {
    const emptyClient = makeMockClient({
      retrieve: vi.fn().mockResolvedValue([]),
    });
    const [tool] = createKnowledgeBaseTools(emptyClient);

    const output = await tool.handler(
      { query: 'nonexistent topic xyz' },
      authorisedCtx
    );

    expect(output.count).toBe(0);
    expect(output.results).toEqual([]);
  });

  // -------------------------------------------------------------------------
  // Scope enforcement
  // -------------------------------------------------------------------------

  it('throws ScopeError when the caller has no scopes', async () => {
    const [tool] = createKnowledgeBaseTools(mockClient);

    await expect(
      tool.handler({ query: 'anything' }, unauthorisedCtx)
    ).rejects.toThrow();

    // The client must NOT be called when auth fails.
    expect(mockClient.retrieve).not.toHaveBeenCalled();
  });

  it('throws ScopeError when the caller only has a finance scope (not ops:read)', async () => {
    const [tool] = createKnowledgeBaseTools(mockClient);

    await expect(
      tool.handler({ query: 'anything' }, financeOnlyCtx)
    ).rejects.toThrow();

    expect(mockClient.retrieve).not.toHaveBeenCalled();
  });

  it('succeeds when the caller has ops:read among multiple scopes', async () => {
    const multiScopeCtx: AuthContext = {
      sub: 'adam@seahavenind.com',
      scopes: ['ops:read', 'ops:tasks', 'finance:read', 'finance:admin'],
      aud: 'sh-mcp-ops',
    };
    const [tool] = createKnowledgeBaseTools(mockClient);

    const output = await tool.handler({ query: 'anything' }, multiScopeCtx);
    expect(output.count).toBe(2);
  });

  // -------------------------------------------------------------------------
  // Client error
  // -------------------------------------------------------------------------

  it('surfaces a client error as a rejected promise', async () => {
    const errorClient = makeMockClient({
      retrieve: vi.fn().mockRejectedValue(new Error('Bedrock Retrieve failed')),
    });
    const [tool] = createKnowledgeBaseTools(errorClient);

    await expect(
      tool.handler({ query: 'HVAC' }, authorisedCtx)
    ).rejects.toThrow('Bedrock Retrieve failed');
  });

  it('surfaces an unexpected error type without swallowing it', async () => {
    const weirdClient = makeMockClient({
      retrieve: vi.fn().mockRejectedValue('string error'),
    });
    const [tool] = createKnowledgeBaseTools(weirdClient);

    await expect(
      tool.handler({ query: 'test' }, authorisedCtx)
    ).rejects.toBe('string error');
  });

  // -------------------------------------------------------------------------
  // Throttle / retry
  // -------------------------------------------------------------------------

  it('propagates a ThrottlingException from the client', async () => {
    // Simulate the shape Bedrock SDK throws for throttling.
    const throttleError = Object.assign(new Error('Too many requests'), {
      name: 'ThrottlingException',
      $fault: 'client',
      $retryable: { throttling: true },
    });

    const throttledClient = makeMockClient({
      retrieve: vi.fn().mockRejectedValue(throttleError),
    });
    const [tool] = createKnowledgeBaseTools(throttledClient);

    const rejection = await tool
      .handler({ query: 'anything' }, authorisedCtx)
      .catch((e: unknown) => e);

    expect((rejection as Error).name).toBe('ThrottlingException');
  });

  it('propagates throttle on first call (retry logic placeholder)', async () => {
    // When retry logic is added (e.g. exponential back-off wrapper), update
    // this test to assert the mock is called N times and eventually succeeds.
    // For now assert the error propagates unchanged so the server layer can
    // apply its own retry strategy.
    const throttleError = Object.assign(new Error('Too many requests'), {
      name: 'ThrottlingException',
    });

    const client = makeMockClient({
      retrieve: vi.fn().mockRejectedValue(throttleError),
    });
    const [tool] = createKnowledgeBaseTools(client);

    await expect(
      tool.handler({ query: 'test' }, authorisedCtx)
    ).rejects.toMatchObject({ name: 'ThrottlingException' });

    // Exactly one attempt — no retry implemented yet.
    expect(client.retrieve).toHaveBeenCalledOnce();
  });

  // -------------------------------------------------------------------------
  // Input schema assertions (contract)
  // -------------------------------------------------------------------------

  it('declares query as a required string in the inputSchema', () => {
    const [tool] = createKnowledgeBaseTools(mockClient);
    const schema = tool.inputSchema as {
      required: string[];
      properties: Record<string, { type: string }>;
    };

    expect(schema.required).toContain('query');
    expect(schema.properties['query'].type).toBe('string');
  });

  it('declares maxResults as an optional integer in the inputSchema', () => {
    const [tool] = createKnowledgeBaseTools(mockClient);
    const schema = tool.inputSchema as {
      required: string[];
      properties: Record<string, { type: string; minimum: number; maximum: number }>;
    };

    expect(schema.required).not.toContain('maxResults');
    expect(schema.properties['maxResults'].type).toBe('integer');
    expect(schema.properties['maxResults'].minimum).toBe(1);
    expect(schema.properties['maxResults'].maximum).toBe(20);
  });
});