mcpctl/src/mcplocal/tests/providers.test.ts

import { describe, it, expect, vi, beforeEach } from 'vitest';
import { ProviderRegistry } from '../src/providers/registry.js';
import type { LlmProvider, CompletionOptions, CompletionResult } from '../src/providers/types.js';

function mockProvider(name: string): LlmProvider {
  return {
    name,
    complete: vi.fn(async (): Promise<CompletionResult> => ({
      content: `Response from ${name}`,
      toolCalls: [],
      usage: { promptTokens: 10, completionTokens: 20, totalTokens: 30 },
      finishReason: 'stop',
    })),
    listModels: vi.fn(async () => [`${name}-model-1`, `${name}-model-2`]),
    isAvailable: vi.fn(async () => true),
  };
}

describe('ProviderRegistry', () => {
  let registry: ProviderRegistry;

  beforeEach(() => {
    registry = new ProviderRegistry();
  });

  it('starts with no providers', () => {
    expect(registry.list()).toEqual([]);
    expect(registry.getActive()).toBeNull();
    expect(registry.getActiveName()).toBeNull();
  });

  it('registers a provider and sets it as active', () => {
    const openai = mockProvider('openai');
    registry.register(openai);

    expect(registry.list()).toEqual(['openai']);
    expect(registry.getActive()).toBe(openai);
    expect(registry.getActiveName()).toBe('openai');
  });

  it('first registered provider becomes active', () => {
    registry.register(mockProvider('openai'));
    registry.register(mockProvider('anthropic'));

    expect(registry.getActiveName()).toBe('openai');
    expect(registry.list()).toEqual(['openai', 'anthropic']);
  });

  it('switches active provider', () => {
    registry.register(mockProvider('openai'));
    registry.register(mockProvider('anthropic'));

    registry.setActive('anthropic');
    expect(registry.getActiveName()).toBe('anthropic');
  });

  it('throws when setting unknown provider as active', () => {
    expect(() => registry.setActive('unknown')).toThrow("Provider 'unknown' is not registered");
  });

  it('gets provider by name', () => {
    const openai = mockProvider('openai');
    registry.register(openai);

    expect(registry.get('openai')).toBe(openai);
    expect(registry.get('unknown')).toBeUndefined();
  });

  it('unregisters a provider', () => {
    registry.register(mockProvider('openai'));
    registry.register(mockProvider('anthropic'));

    registry.unregister('openai');
    expect(registry.list()).toEqual(['anthropic']);
    // Active should switch to remaining provider
    expect(registry.getActiveName()).toBe('anthropic');
  });

  it('unregistering active provider switches to next available', () => {
    registry.register(mockProvider('openai'));
    registry.register(mockProvider('anthropic'));
    registry.setActive('openai');

    registry.unregister('openai');
    expect(registry.getActiveName()).toBe('anthropic');
  });

  it('unregistering last provider clears active', () => {
    registry.register(mockProvider('openai'));
    registry.unregister('openai');

    expect(registry.getActive()).toBeNull();
    expect(registry.getActiveName()).toBeNull();
  });

  it('active provider can complete', async () => {
    const provider = mockProvider('openai');
    registry.register(provider);

    const active = registry.getActive()!;
    const result = await active.complete({
      messages: [{ role: 'user', content: 'Hello' }],
    });

    expect(result.content).toBe('Response from openai');
    expect(result.finishReason).toBe('stop');
    expect(provider.complete).toHaveBeenCalled();
  });

  it('active provider can list models', async () => {
    registry.register(mockProvider('anthropic'));

    const active = registry.getActive()!;
    const models = await active.listModels();

    expect(models).toEqual(['anthropic-model-1', 'anthropic-model-2']);
  });

  describe('tier management', () => {
    it('assigns providers to tiers', () => {
      registry.register(mockProvider('vllm'));
      registry.register(mockProvider('gemini'));

      registry.assignTier('vllm', 'fast');
      registry.assignTier('gemini', 'heavy');

      expect(registry.getTierProviders('fast')).toEqual(['vllm']);
      expect(registry.getTierProviders('heavy')).toEqual(['gemini']);
      expect(registry.hasTierConfig()).toBe(true);
    });

    it('getProvider returns tier-specific provider', () => {
      const vllm = mockProvider('vllm');
      const gemini = mockProvider('gemini');
      registry.register(vllm);
      registry.register(gemini);

      registry.assignTier('vllm', 'fast');
      registry.assignTier('gemini', 'heavy');

      expect(registry.getProvider('fast')).toBe(vllm);
      expect(registry.getProvider('heavy')).toBe(gemini);
    });

    it('getProvider falls back to other tier', () => {
      const vllm = mockProvider('vllm');
      registry.register(vllm);

      registry.assignTier('vllm', 'fast');

      // Requesting heavy but only fast exists → falls back to fast
      expect(registry.getProvider('heavy')).toBe(vllm);
    });

    it('getProvider falls back to getActive when no tiers', () => {
      const openai = mockProvider('openai');
      registry.register(openai);

      // No tier assignments → falls back to legacy getActive()
      expect(registry.getProvider('fast')).toBe(openai);
      expect(registry.getProvider('heavy')).toBe(openai);
      expect(registry.hasTierConfig()).toBe(false);
    });

    it('unregister removes from tier assignments', () => {
      registry.register(mockProvider('vllm'));
      registry.register(mockProvider('gemini'));

      registry.assignTier('vllm', 'fast');
      registry.assignTier('gemini', 'heavy');

      registry.unregister('vllm');

      expect(registry.getTierProviders('fast')).toEqual([]);
      expect(registry.getTierProviders('heavy')).toEqual(['gemini']);
    });

    it('assignTier throws for unregistered provider', () => {
      expect(() => registry.assignTier('unknown', 'fast')).toThrow("Provider 'unknown' is not registered");
    });

    it('multiple providers in same tier uses first', () => {
      const vllm = mockProvider('vllm');
      const ollama = mockProvider('ollama');
      registry.register(vllm);
      registry.register(ollama);

      registry.assignTier('vllm', 'fast');
      registry.assignTier('ollama', 'fast');

      expect(registry.getProvider('fast')).toBe(vllm);
      expect(registry.getTierProviders('fast')).toEqual(['vllm', 'ollama']);
    });

    it('listProviders includes tier info', () => {
      registry.register(mockProvider('vllm'));
      registry.register(mockProvider('gemini'));

      registry.assignTier('vllm', 'fast');
      registry.assignTier('gemini', 'heavy');

      const providers = registry.listProviders();
      expect(providers).toEqual([
        { name: 'vllm', tiers: ['fast'] },
        { name: 'gemini', tiers: ['heavy'] },
      ]);
    });

    it('disposeAll calls dispose on all providers', () => {
      const disposeFn = vi.fn();
      const provider = { ...mockProvider('test'), dispose: disposeFn };
      registry.register(provider);

      registry.disposeAll();

      expect(disposeFn).toHaveBeenCalledOnce();
    });
  });
});