dify/web/app/components/header/account-setting/model-provider-page/utils.spec.ts

import type { Mock } from 'vitest'
import { describe, expect, it, vi } from 'vitest'
import {
  deleteModelProvider,
  setModelProvider,
  validateModelLoadBalancingCredentials,
  validateModelProvider,
} from '@/service/common'
import { ValidatedStatus } from '../key-validator/declarations'
import {
  ConfigurationMethodEnum,
  FormTypeEnum,
  ModelTypeEnum,
} from './declarations'
import {
  genModelNameFormSchema,
  genModelTypeFormSchema,
  modelTypeFormat,
  removeCredentials,
  saveCredentials,
  savePredefinedLoadBalancingConfig,
  sizeFormat,
  validateCredentials,
  validateLoadBalancingCredentials,
} from './utils'

// Mock service/common functions
vi.mock('@/service/common', () => ({
  deleteModelProvider: vi.fn(),
  setModelProvider: vi.fn(),
  validateModelLoadBalancingCredentials: vi.fn(),
  validateModelProvider: vi.fn(),
}))

describe('utils', () => {
  afterEach(() => {
    vi.clearAllMocks()
  })

  describe('sizeFormat', () => {
    it('should format size less than 1000', () => {
      expect(sizeFormat(500)).toBe('500')
    })

    it('should format size greater than 1000', () => {
      expect(sizeFormat(1500)).toBe('1K')
    })
  })

  describe('modelTypeFormat', () => {
    it('should format text embedding type', () => {
      expect(modelTypeFormat(ModelTypeEnum.textEmbedding)).toBe('TEXT EMBEDDING')
    })

    it('should format other types', () => {
      expect(modelTypeFormat(ModelTypeEnum.textGeneration)).toBe('LLM')
    })
  })

  describe('validateCredentials', () => {
    it('should validate predefined credentials successfully', async () => {
      (validateModelProvider as unknown as Mock).mockResolvedValue({ result: 'success' })
      const result = await validateCredentials(true, 'provider', { key: 'value' })
      expect(result).toEqual({ status: ValidatedStatus.Success })
      expect(validateModelProvider).toHaveBeenCalledWith({
        url: '/workspaces/current/model-providers/provider/credentials/validate',
        body: { credentials: { key: 'value' } },
      })
    })

    it('should validate custom credentials successfully', async () => {
      (validateModelProvider as unknown as Mock).mockResolvedValue({ result: 'success' })
      const result = await validateCredentials(false, 'provider', {
        __model_name: 'model',
        __model_type: 'type',
        key: 'value',
      })
      expect(result).toEqual({ status: ValidatedStatus.Success })
      expect(validateModelProvider).toHaveBeenCalledWith({
        url: '/workspaces/current/model-providers/provider/models/credentials/validate',
        body: {
          model: 'model',
          model_type: 'type',
          credentials: { key: 'value' },
        },
      })
    })

    it('should handle validation failure', async () => {
      (validateModelProvider as unknown as Mock).mockResolvedValue({ result: 'error', error: 'failed' })
      const result = await validateCredentials(true, 'provider', {})
      expect(result).toEqual({ status: ValidatedStatus.Error, message: 'failed' })
    })

    it('should handle exception', async () => {
      (validateModelProvider as unknown as Mock).mockRejectedValue(new Error('network error'))
      const result = await validateCredentials(true, 'provider', {})
      expect(result).toEqual({ status: ValidatedStatus.Error, message: 'network error' })
    })
  })

  describe('validateLoadBalancingCredentials', () => {
    it('should validate load balancing credentials successfully', async () => {
      (validateModelLoadBalancingCredentials as unknown as Mock).mockResolvedValue({ result: 'success' })
      const result = await validateLoadBalancingCredentials(true, 'provider', {
        __model_name: 'model',
        __model_type: 'type',
        key: 'value',
      })
      expect(result).toEqual({ status: ValidatedStatus.Success })
      expect(validateModelLoadBalancingCredentials).toHaveBeenCalledWith({
        url: '/workspaces/current/model-providers/provider/models/load-balancing-configs/credentials-validate',
        body: {
          model: 'model',
          model_type: 'type',
          credentials: { key: 'value' },
        },
      })
    })
    it('should validate load balancing credentials successfully with id', async () => {
      (validateModelLoadBalancingCredentials as unknown as Mock).mockResolvedValue({ result: 'success' })
      const result = await validateLoadBalancingCredentials(true, 'provider', {
        __model_name: 'model',
        __model_type: 'type',
        key: 'value',
      }, 'id')
      expect(result).toEqual({ status: ValidatedStatus.Success })
      expect(validateModelLoadBalancingCredentials).toHaveBeenCalledWith({
        url: '/workspaces/current/model-providers/provider/models/load-balancing-configs/id/credentials-validate',
        body: {
          model: 'model',
          model_type: 'type',
          credentials: { key: 'value' },
        },
      })
    })

    it('should handle validation failure', async () => {
      (validateModelLoadBalancingCredentials as unknown as Mock).mockResolvedValue({ result: 'error', error: 'failed' })
      const result = await validateLoadBalancingCredentials(true, 'provider', {})
      expect(result).toEqual({ status: ValidatedStatus.Error, message: 'failed' })
    })
  })

  describe('saveCredentials', () => {
    it('should save predefined credentials', async () => {
      await saveCredentials(true, 'provider', { __authorization_name__: 'name', key: 'value' })
      expect(setModelProvider).toHaveBeenCalledWith({
        url: '/workspaces/current/model-providers/provider/credentials',
        body: {
          config_from: ConfigurationMethodEnum.predefinedModel,
          credentials: { key: 'value' },
          load_balancing: undefined,
          name: 'name',
        },
      })
    })

    it('should save custom credentials', async () => {
      await saveCredentials(false, 'provider', {
        __model_name: 'model',
        __model_type: 'type',
        key: 'value',
      })
      expect(setModelProvider).toHaveBeenCalledWith({
        url: '/workspaces/current/model-providers/provider/models',
        body: {
          model: 'model',
          model_type: 'type',
          credentials: { key: 'value' },
          load_balancing: undefined,
        },
      })
    })
  })

  describe('savePredefinedLoadBalancingConfig', () => {
    it('should save predefined load balancing config', async () => {
      await savePredefinedLoadBalancingConfig('provider', {
        __model_name: 'model',
        __model_type: 'type',
        key: 'value',
      })
      expect(setModelProvider).toHaveBeenCalledWith({
        url: '/workspaces/current/model-providers/provider/models',
        body: {
          config_from: ConfigurationMethodEnum.predefinedModel,
          model: 'model',
          model_type: 'type',
          credentials: { key: 'value' },
          load_balancing: undefined,
        },
      })
    })
  })

  describe('removeCredentials', () => {
    it('should remove predefined credentials', async () => {
      await removeCredentials(true, 'provider', {}, 'id')
      expect(deleteModelProvider).toHaveBeenCalledWith({
        url: '/workspaces/current/model-providers/provider/credentials',
        body: { credential_id: 'id' },
      })
    })

    it('should remove custom credentials', async () => {
      await removeCredentials(false, 'provider', {
        __model_name: 'model',
        __model_type: 'type',
      })
      expect(deleteModelProvider).toHaveBeenCalledWith({
        url: '/workspaces/current/model-providers/provider/models',
        body: {
          model: 'model',
          model_type: 'type',
        },
      })
    })
  })

  describe('genModelTypeFormSchema', () => {
    it('should generate form schema', () => {
      const schema = genModelTypeFormSchema([ModelTypeEnum.textGeneration])
      expect(schema.type).toBe(FormTypeEnum.select)
      expect(schema.variable).toBe('__model_type')
      expect(schema.options[0].value).toBe(ModelTypeEnum.textGeneration)
    })
  })

  describe('genModelNameFormSchema', () => {
    it('should generate form schema', () => {
      const schema = genModelNameFormSchema()
      expect(schema.type).toBe(FormTypeEnum.textInput)
      expect(schema.variable).toBe('__model_name')
      expect(schema.required).toBe(true)
    })
  })
})