learning_ai_common_plat/packages/ollama-client/src/client.test.ts

import { describe, it, expect, vi, beforeEach } from 'vitest';
import { OllamaClient } from './client.js';

const BASE_URL = 'http://localhost:11434';

function mockFetch(response: unknown, options?: { ok?: boolean; status?: number }) {
  return vi.fn().mockResolvedValue({
    ok: options?.ok ?? true,
    status: options?.status ?? 200,
    json: () => Promise.resolve(response),
    text: () => Promise.resolve(JSON.stringify(response)),
  });
}

describe('OllamaClient', () => {
  let client: OllamaClient;

  beforeEach(() => {
    client = new OllamaClient({ baseUrl: BASE_URL });
    vi.restoreAllMocks();
  });

  describe('constructor', () => {
    it('strips trailing slashes from baseUrl', () => {
      const c = new OllamaClient({ baseUrl: 'http://localhost:11434///' });
      expect(c.baseUrl).toBe('http://localhost:11434');
    });
  });

  describe('tags', () => {
    it('returns list of models', async () => {
      const models = [{ name: 'llama3', size: 1000, digest: 'abc', modified_at: '2024-01-01' }];
      globalThis.fetch = mockFetch({ models });

      const result = await client.tags();
      expect(result).toEqual(models);
      expect(globalThis.fetch).toHaveBeenCalledWith(
        `${BASE_URL}/api/tags`,
        expect.objectContaining({
          headers: expect.objectContaining({ 'Content-Type': 'application/json' }),
        })
      );
    });

    it('returns empty array when models is null', async () => {
      globalThis.fetch = mockFetch({ models: null });
      const result = await client.tags();
      expect(result).toEqual([]);
    });
  });

  describe('ps', () => {
    it('returns running models', async () => {
      const models = [
        { name: 'llama3', size: 1000, digest: 'abc', expires_at: '2024-01-01', size_vram: 500 },
      ];
      globalThis.fetch = mockFetch({ models });

      const result = await client.ps();
      expect(result).toEqual(models);
    });
  });

  describe('show', () => {
    it('sends POST with model name', async () => {
      const showData = { modelfile: '', parameters: '', template: '', details: {} };
      globalThis.fetch = mockFetch(showData);

      const result = await client.show('llama3');
      expect(result).toEqual(showData);
      expect(globalThis.fetch).toHaveBeenCalledWith(
        `${BASE_URL}/api/show`,
        expect.objectContaining({
          method: 'POST',
          body: JSON.stringify({ name: 'llama3' }),
        })
      );
    });
  });

  describe('pull (non-streaming)', () => {
    it('pulls a model', async () => {
      globalThis.fetch = mockFetch({ status: 'success' });

      const result = await client.pull('llama3');
      expect(result).toEqual({ status: 'success' });
    });

    it('throws on failure', async () => {
      globalThis.fetch = mockFetch('Pull failed', { ok: false, status: 500 });
      await expect(client.pull('bad-model')).rejects.toThrow('Ollama pull failed (500)');
    });
  });

  describe('load', () => {
    it('sends generate with keep_alive', async () => {
      globalThis.fetch = mockFetch({});

      await client.load('llama3', '15m');
      expect(globalThis.fetch).toHaveBeenCalledWith(
        `${BASE_URL}/api/generate`,
        expect.objectContaining({
          method: 'POST',
          body: JSON.stringify({ model: 'llama3', prompt: '', keep_alive: '15m' }),
        })
      );
    });
  });

  describe('unload', () => {
    it('sends generate with keep_alive: 0', async () => {
      globalThis.fetch = mockFetch({});

      await client.unload('llama3');
      expect(globalThis.fetch).toHaveBeenCalledWith(
        `${BASE_URL}/api/generate`,
        expect.objectContaining({
          method: 'POST',
          body: JSON.stringify({ model: 'llama3', prompt: '', keep_alive: '0' }),
        })
      );
    });
  });

  describe('delete', () => {
    it('sends DELETE request', async () => {
      globalThis.fetch = vi.fn().mockResolvedValue({ ok: true, status: 200 });

      await client.delete('llama3');
      expect(globalThis.fetch).toHaveBeenCalledWith(
        `${BASE_URL}/api/delete`,
        expect.objectContaining({
          method: 'DELETE',
          body: JSON.stringify({ name: 'llama3' }),
        })
      );
    });

    it('throws on failure', async () => {
      globalThis.fetch = vi.fn().mockResolvedValue({
        ok: false,
        status: 404,
        text: () => Promise.resolve('model not found'),
      });
      await expect(client.delete('nope')).rejects.toThrow('Ollama delete failed (404)');
    });
  });

  describe('version', () => {
    it('returns version string', async () => {
      globalThis.fetch = mockFetch({ version: '0.5.4' });

      const result = await client.version();
      expect(result).toBe('0.5.4');
    });
  });
});