brainy/tests/custom-models-path.test.ts

/**
 * Custom Models Path Test
 * 
 * Tests the custom models directory functionality for Docker deployments
 */

import { describe, it, expect, beforeAll, afterAll, vi } from 'vitest'
import { RobustModelLoader } from '../src/utils/robustModelLoader.js'
import { writeFile, mkdir, rm } from 'fs/promises'
import { join } from 'path'
import { tmpdir } from 'os'

describe('Custom Models Path', () => {
  let tempDir: string
  let originalEnv: string | undefined

  beforeAll(() => {
    // Save original environment variable
    originalEnv = process.env.BRAINY_MODELS_PATH
  })

  afterAll(() => {
    // Restore original environment variable
    if (originalEnv !== undefined) {
      process.env.BRAINY_MODELS_PATH = originalEnv
    } else {
      delete process.env.BRAINY_MODELS_PATH
    }
  })

  beforeEach(async () => {
    // Create temporary directory for each test
    tempDir = join(tmpdir(), `brainy-test-${Date.now()}`)
    await mkdir(tempDir, { recursive: true })
  })

  afterEach(async () => {
    // Clean up temporary directory
    try {
      await rm(tempDir, { recursive: true, force: true })
    } catch (error) {
      console.error('Cleanup error:', error)
    }
  })

  it('should use BRAINY_MODELS_PATH environment variable', async () => {
    // Set environment variable
    process.env.BRAINY_MODELS_PATH = tempDir

    const loader = new RobustModelLoader({ verbose: true })
    expect((loader as any).options.customModelsPath).toBe(tempDir)
  })

  it('should use MODELS_PATH environment variable as fallback', async () => {
    // Clear BRAINY_MODELS_PATH and set MODELS_PATH
    delete process.env.BRAINY_MODELS_PATH
    process.env.MODELS_PATH = tempDir

    const loader = new RobustModelLoader({ verbose: true })
    expect((loader as any).options.customModelsPath).toBe(tempDir)

    // Clean up
    delete process.env.MODELS_PATH
  })

  it('should prioritize customModelsPath option over environment variables', async () => {
    process.env.BRAINY_MODELS_PATH = '/env/path'
    const customPath = '/custom/path'

    const loader = new RobustModelLoader({ 
      customModelsPath: customPath,
      verbose: true 
    })
    
    expect((loader as any).options.customModelsPath).toBe(customPath)
  })

  it('should check multiple subdirectories for models', async () => {
    const loader = new RobustModelLoader({ 
      customModelsPath: tempDir,
      verbose: true 
    })

    // Create a mock model.json in one of the expected subdirectories
    const modelDir = join(tempDir, 'universal-sentence-encoder')
    await mkdir(modelDir, { recursive: true })
    
    const mockModelJson = {
      format: 'tfjs-graph-model',
      modelTopology: {},
      weightsManifest: []
    }
    
    await writeFile(
      join(modelDir, 'model.json'), 
      JSON.stringify(mockModelJson, null, 2)
    )

    // Mock the tryLoadFromCustomPath method to avoid actual TensorFlow loading
    const tryLoadSpy = vi.spyOn(loader as any, 'tryLoadFromCustomPath')
    tryLoadSpy.mockResolvedValue(null) // Return null to avoid complex mocking

    await (loader as any).tryLoadLocalBundledModel()

    // Verify the method was called with the correct path
    expect(tryLoadSpy).toHaveBeenCalledWith(tempDir)
  })

  it('should log helpful messages when checking custom path', async () => {
    const consoleSpy = vi.spyOn(console, 'log').mockImplementation(() => {})
    
    const loader = new RobustModelLoader({ 
      customModelsPath: tempDir,
      verbose: true 
    })

    try {
      await (loader as any).tryLoadLocalBundledModel()
    } catch (error) {
      // Expected in test environment
    }

    // Check that it logged the custom path check
    expect(consoleSpy).toHaveBeenCalledWith(
      expect.stringContaining(`Checking custom models directory: ${tempDir}`)
    )

    consoleSpy.mockRestore()
  })

  it('should handle non-existent custom paths gracefully', async () => {
    const nonExistentPath = '/this/path/does/not/exist'
    const consoleSpy = vi.spyOn(console, 'log').mockImplementation(() => {})
    
    const loader = new RobustModelLoader({ 
      customModelsPath: nonExistentPath,
      verbose: true 
    })

    const result = await (loader as any).tryLoadFromCustomPath(nonExistentPath)
    expect(result).toBeNull()

    // Should log that no model was found
    expect(consoleSpy).toHaveBeenCalledWith(
      expect.stringContaining(`No model found in custom path: ${nonExistentPath}`)
    )

    consoleSpy.mockRestore()
  })

  it('should provide helpful warning messages about custom paths', async () => {
    const consoleSpy = vi.spyOn(console, 'warn').mockImplementation(() => {})
    
    // Test without any custom path set
    const loader = new RobustModelLoader({ verbose: false })

    try {
      await loader.loadModelWithFallbacks()
    } catch (error) {
      // Expected in test environment without actual models
    }

    // Check that the warning mentions the custom path option
    const warnCalls = consoleSpy.mock.calls.flat()
    const hasCustomPathMention = warnCalls.some(call => 
      typeof call === 'string' && call.includes('BRAINY_MODELS_PATH')
    )

    expect(hasCustomPathMention).toBe(true)

    consoleSpy.mockRestore()
  })
})