promptfoo--promptfoo/test/providers/fal.test.ts

import { afterEach, beforeEach, describe, expect, it, vi } from 'vitest';
import { getCache, isCacheEnabled } from '../../src/cache';
import { getEnvString } from '../../src/envars';
import { FalImageGenerationProvider } from '../../src/providers/fal';

const mockSubscribe = vi.hoisted(() => vi.fn());
const mockConfig = vi.hoisted(() => vi.fn());

vi.mock('@fal-ai/client', async (importOriginal) => {
  return {
    ...(await importOriginal()),

    fal: {
      subscribe: mockSubscribe,
      config: mockConfig,
    },
  };
});

vi.mock('../../src/cache', async () => {
  const actual = await vi.importActual<typeof import('../../src/cache')>('../../src/cache');
  return {
    ...actual,
    getCache: vi.fn(),
    isCacheEnabled: vi.fn(),
  };
});

vi.mock('../../src/envars', async (importOriginal) => {
  return {
    ...(await importOriginal()),
    getEnvString: vi.fn(),
    getEnvInt: vi.fn().mockReturnValue(300000),
    getEnvBool: vi.fn().mockReturnValue(true),
    getEnvFloat: vi.fn(),
  };
});

describe('Fal Provider', () => {
  beforeEach(() => {
    vi.clearAllMocks();
    vi.mocked(isCacheEnabled).mockReturnValue(false);
    vi.mocked(getCache).mockReturnValue({
      get: vi.fn().mockResolvedValue(null),
      set: vi.fn(),
      wrap: vi.fn(),
      del: vi.fn(),
      clear: vi.fn(),
      stores: [
        {
          get: vi.fn(),
          set: vi.fn(),
        },
      ] as any,
      mget: vi.fn(),
      mset: vi.fn(),
      mdel: vi.fn(),
      reset: vi.fn(),
      ttl: vi.fn(),
      on: vi.fn(),
      removeAllListeners: vi.fn(),
    } as any);
  });

  afterEach(() => {
    vi.resetAllMocks();
  });

  describe('FalImageGenerationProvider', () => {
    let provider: FalImageGenerationProvider;

    beforeEach(() => {
      provider = new FalImageGenerationProvider('fal-ai/flux/schnell', {
        config: { apiKey: 'test-api-key' },
      });
    });

    describe('constructor and configuration', () => {
      it('should create provider with default options', () => {
        expect(provider).toEqual(
          expect.objectContaining({
            modelName: 'fal-ai/flux/schnell',
            modelType: 'image',
            apiKey: 'test-api-key',
          }),
        );
      });

      it('should create provider with custom configuration options', () => {
        const customProvider = new FalImageGenerationProvider('fal-ai/fast-sdxl', {
          config: {
            apiKey: 'custom-key',
            seed: 12345,
            num_inference_steps: 8,
            image_size: { width: 1024, height: 1024 },
          },
        });

        expect(customProvider).toEqual(
          expect.objectContaining({
            modelName: 'fal-ai/fast-sdxl',
            apiKey: 'custom-key',
            config: {
              apiKey: 'custom-key',
              seed: 12345,
              num_inference_steps: 8,
              image_size: { width: 1024, height: 1024 },
            },
          }),
        );
      });

      it('should use environment variable for API key when not provided in config', () => {
        vi.mocked(getEnvString).mockImplementation(function () {
          return 'env-api-key';
        });

        const envProvider = new FalImageGenerationProvider('fal-ai/flux/schnell');

        expect(envProvider.apiKey).toBe('env-api-key');
        expect(getEnvString).toHaveBeenCalledWith('FAL_KEY');
      });

      it('should use env override for API key', () => {
        const envProvider = new FalImageGenerationProvider('fal-ai/flux/schnell', {
          env: { FAL_KEY: 'override-key' },
        });

        expect(envProvider.apiKey).toBe('override-key');
      });
    });

    describe('provider identification', () => {
      it('should generate correct provider id', () => {
        expect(provider.id()).toBe('fal:image:fal-ai/flux/schnell');
      });

      it('should generate correct string representation', () => {
        expect(provider.toString()).toBe('[fal.ai Image Generation Provider fal-ai/flux/schnell]');
      });
    });

    describe('API key validation', () => {
      it('should throw error when API key is not set', async () => {
        vi.mocked(getEnvString).mockImplementation(function () {
          return undefined as any;
        });

        const noKeyProvider = new FalImageGenerationProvider('fal-ai/flux/schnell');

        await expect(noKeyProvider.callApi('test prompt')).rejects.toThrow(
          'fal.ai API key is not set. Set the FAL_KEY environment variable or or add `apiKey` to the provider config.',
        );
      });
    });

    describe('API calls and image generation', () => {
      const mockImageResponse = {
        data: {
          images: [{ url: 'https://example.com/image.png' }],
        },
        requestId: 'test-request-id',
      };

      const mockSingleImageResponse = {
        data: {
          image: { url: 'https://example.com/image.png' },
        },
        requestId: 'test-request-id',
      };

      it('should call fal API and return markdown image with images array response', async () => {
        mockSubscribe.mockResolvedValueOnce(mockImageResponse);

        const result = await provider.callApi('a cute cat');

        expect(mockConfig).toHaveBeenCalledWith({
          credentials: 'test-api-key',
        });
        expect(mockSubscribe).toHaveBeenCalledWith('fal-ai/flux/schnell', {
          input: {
            prompt: 'a cute cat',
          },
        });
        expect(result).toEqual({
          cached: false,
          output: '![a cute cat](https://example.com/image.png)',
        });
      });

      it('should call fal API and return markdown image with single image response', async () => {
        mockSubscribe.mockResolvedValueOnce(mockSingleImageResponse);

        const result = await provider.callApi('a beautiful landscape');

        expect(result).toEqual({
          cached: false,
          output: '![a beautiful landscape](https://example.com/image.png)',
        });
      });

      it('should merge provider config with prompt context config', async () => {
        const providerWithConfig = new FalImageGenerationProvider('fal-ai/flux/schnell', {
          config: {
            apiKey: 'test-api-key',
            seed: 12345,
            num_inference_steps: 8,
          },
        });

        mockSubscribe.mockResolvedValueOnce(mockImageResponse);

        await providerWithConfig.callApi('test prompt', {
          prompt: {
            raw: 'test prompt',
            label: 'test',
            config: {
              guidance_scale: 7.5,
            },
          },
          vars: {},
        });

        expect(mockSubscribe).toHaveBeenCalledWith('fal-ai/flux/schnell', {
          input: {
            prompt: 'test prompt',
            seed: 12345,
            num_inference_steps: 8,
            guidance_scale: 7.5,
          },
        });
      });

      it('should pass fal client proxy config without sending it as model input', async () => {
        const providerWithClientConfig = new FalImageGenerationProvider('fal-ai/flux/schnell', {
          config: {
            apiKey: 'test-api-key',
            seed: 12345,
            client: {
              proxyUrl: {
                url: 'http://fal-proxy.test/api/fal',
                when: 'always',
              },
            },
          },
        });

        mockSubscribe.mockResolvedValueOnce(mockImageResponse);

        await providerWithClientConfig.callApi('test prompt', {
          prompt: {
            raw: 'test prompt',
            label: 'test',
            config: {
              guidance_scale: 7.5,
              client: {
                proxyUrl: 'http://context-proxy.test/api/fal',
              },
            },
          },
          vars: {},
        });

        expect(mockConfig).toHaveBeenCalledWith({
          credentials: 'test-api-key',
          proxyUrl: {
            url: 'http://fal-proxy.test/api/fal',
            when: 'always',
          },
        });
        expect(mockSubscribe).toHaveBeenCalledWith('fal-ai/flux/schnell', {
          input: {
            prompt: 'test prompt',
            seed: 12345,
            guidance_scale: 7.5,
          },
        });
      });
    });

    describe('runInference method', () => {
      it('should resolve image URL from images array', async () => {
        const mockResponse = {
          data: {
            images: [{ url: 'https://example.com/image.png' }],
          },
          requestId: 'test-request-id',
        };
        mockSubscribe.mockResolvedValueOnce(mockResponse);

        const result = await provider.runInference({
          prompt: 'a cute cat',
          seed: 12345,
        });

        expect(result).toBe('![a cute cat](https://example.com/image.png)');
      });

      it('should resolve image URL from single image object', async () => {
        const mockResponse = {
          data: {
            image: { url: 'https://example.com/image.png' },
          },
          requestId: 'test-request-id',
        };
        mockSubscribe.mockResolvedValueOnce(mockResponse);

        const result = await provider.runInference({
          prompt: 'a beautiful landscape',
          seed: 12345,
        });

        expect(result).toBe('![a beautiful landscape](https://example.com/image.png)');
      });
    });

    describe('prompt processing', () => {
      const mockResponse = {
        data: {
          images: [{ url: 'https://example.com/image.png' }],
        },
        requestId: 'test-request-id',
      };

      beforeEach(() => {
        mockSubscribe.mockResolvedValue(mockResponse);
      });

      it('should sanitize prompt in markdown output', async () => {
        const result = await provider.runInference({
          prompt: 'a [test] prompt\nwith newlines\rand brackets',
        });

        expect(result).toBe(
          '![a (test) prompt with newlines and brackets](https://example.com/image.png)',
        );
      });

      it('should ellipsize long prompts in markdown output', async () => {
        const longPrompt =
          'a very long prompt that exceeds the maximum length and should be ellipsized';

        const result = await provider.runInference({
          prompt: longPrompt,
        });

        expect(result).toBe(
          '![a very long prompt that exceeds the maximum len...](https://example.com/image.png)',
        );
      });
    });

    describe('caching behavior', () => {
      it('should use cached response when cache is enabled and available', async () => {
        vi.mocked(isCacheEnabled).mockImplementation(function () {
          return true;
        });
        const mockCachedResponse = JSON.stringify(
          '![cached prompt](https://cached.example.com/image.png)',
        );

        const mockCache = {
          get: vi.fn().mockResolvedValue(mockCachedResponse),
          set: vi.fn(),
          wrap: vi.fn(),
          del: vi.fn(),
          clear: vi.fn(),
          stores: [
            {
              get: vi.fn(),
              set: vi.fn(),
            },
          ] as any,
          mget: vi.fn(),
          mset: vi.fn(),
          mdel: vi.fn(),
          reset: vi.fn(),
          ttl: vi.fn(),
          on: vi.fn(),
          removeAllListeners: vi.fn(),
        };
        vi.mocked(getCache).mockReturnValue(mockCache as any);

        const result = await provider.callApi('test prompt');

        expect(result).toEqual({
          cached: true,
          output: '![cached prompt](https://cached.example.com/image.png)',
        });
        expect(mockSubscribe).not.toHaveBeenCalled();
        expect(mockCache.get).toHaveBeenCalledWith(
          expect.stringContaining('fal:fal-ai/flux/schnell:'),
        );
      });

      it('should set cache when enabled and response is fresh', async () => {
        vi.mocked(isCacheEnabled).mockImplementation(function () {
          return true;
        });
        const mockCache = {
          get: vi.fn().mockResolvedValue(null),
          set: vi.fn(),
          wrap: vi.fn(),
          del: vi.fn(),
          clear: vi.fn(),
          stores: [
            {
              get: vi.fn(),
              set: vi.fn(),
            },
          ] as any,
          mget: vi.fn(),
          mset: vi.fn(),
          mdel: vi.fn(),
          reset: vi.fn(),
          ttl: vi.fn(),
          on: vi.fn(),
          removeAllListeners: vi.fn(),
        };
        vi.mocked(getCache).mockReturnValue(mockCache as any);

        const mockResponse = {
          data: {
            images: [{ url: 'https://example.com/image.png' }],
          },
          requestId: 'test-request-id',
        };
        mockSubscribe.mockResolvedValueOnce(mockResponse);

        const result = await provider.callApi('test prompt');

        expect(result.cached).toBe(false);
        expect(mockSubscribe).toHaveBeenCalledWith('fal-ai/flux/schnell', {
          input: {
            prompt: 'test prompt',
          },
        });
        expect(mockCache.set).toHaveBeenCalledWith(
          expect.stringContaining('fal:fal-ai/flux/schnell:'),
          JSON.stringify('![test prompt](https://example.com/image.png)'),
        );
      });

      it('should hash prompt and config values in cache keys', async () => {
        vi.mocked(isCacheEnabled).mockImplementation(function () {
          return true;
        });
        const mockCache = {
          get: vi.fn().mockResolvedValue(null),
          set: vi.fn(),
          wrap: vi.fn(),
          del: vi.fn(),
          clear: vi.fn(),
          stores: [
            {
              get: vi.fn(),
              set: vi.fn(),
            },
          ] as any,
          mget: vi.fn(),
          mset: vi.fn(),
          mdel: vi.fn(),
          reset: vi.fn(),
          ttl: vi.fn(),
          on: vi.fn(),
          removeAllListeners: vi.fn(),
        };
        vi.mocked(getCache).mockReturnValue(mockCache as any);

        mockSubscribe.mockResolvedValueOnce({
          data: {
            images: [{ url: 'https://example.com/image.png' }],
          },
          requestId: 'test-request-id',
        });

        const prompt = 'PFQA_FAL_PROMPT_SENTINEL';
        const contextSecret = 'PFQA_FAL_CONFIG_SECRET_SENTINEL';
        const result = await provider.callApi(prompt, {
          prompt: {
            raw: prompt,
            label: 'test',
            config: {
              negative_prompt: contextSecret,
            },
          },
          vars: {},
        });

        expect(result.cached).toBe(false);
        const cacheKey = mockCache.get.mock.calls[0]?.[0] as string;
        expect(cacheKey).toMatch(
          /^fal:fal-ai\/flux\/schnell:[a-f0-9]{64}:[a-f0-9]{64}:[a-f0-9]{64}$/,
        );
        expect(cacheKey).not.toContain(prompt);
        expect(cacheKey).not.toContain(contextSecret);
        expect(mockCache.set).toHaveBeenCalledWith(
          cacheKey,
          JSON.stringify(`![${prompt}](https://example.com/image.png)`),
        );
      });

      it('should not expose API key values in hashed cache keys', async () => {
        vi.mocked(isCacheEnabled).mockImplementation(function () {
          return true;
        });
        const mockCache = {
          get: vi.fn().mockResolvedValue(null),
          set: vi.fn(),
          wrap: vi.fn(),
          del: vi.fn(),
          clear: vi.fn(),
          stores: [
            {
              get: vi.fn(),
              set: vi.fn(),
            },
          ] as any,
          mget: vi.fn(),
          mset: vi.fn(),
          mdel: vi.fn(),
          reset: vi.fn(),
          ttl: vi.fn(),
          on: vi.fn(),
          removeAllListeners: vi.fn(),
        };
        vi.mocked(getCache).mockReturnValue(mockCache as any);
        mockSubscribe
          .mockResolvedValueOnce({
            data: {
              images: [{ url: 'https://example.com/tenant-a.png' }],
            },
            requestId: 'tenant-a-request-id',
          })
          .mockResolvedValueOnce({
            data: {
              images: [{ url: 'https://example.com/tenant-b.png' }],
            },
            requestId: 'tenant-b-request-id',
          });

        const providerA = new FalImageGenerationProvider('fal-ai/flux/schnell', {
          config: { apiKey: 'fal-tenant-a-secret' },
        });
        const providerB = new FalImageGenerationProvider('fal-ai/flux/schnell', {
          config: { apiKey: 'fal-tenant-b-secret' },
        });

        await providerA.callApi('Shared fal prompt');
        await providerB.callApi('Shared fal prompt');

        const cacheKeyA = mockCache.get.mock.calls[0][0] as string;
        const cacheKeyB = mockCache.get.mock.calls[1][0] as string;
        expect(cacheKeyA).toMatch(
          /^fal:fal-ai\/flux\/schnell:[a-f0-9]{64}:[a-f0-9]{64}:[a-f0-9]{64}$/,
        );
        expect(cacheKeyB).toMatch(
          /^fal:fal-ai\/flux\/schnell:[a-f0-9]{64}:[a-f0-9]{64}:[a-f0-9]{64}$/,
        );
        expect(cacheKeyA).not.toBe(cacheKeyB);
        expect(mockSubscribe).toHaveBeenCalledTimes(2);
        for (const cacheKey of [cacheKeyA, cacheKeyB]) {
          expect(cacheKey).not.toContain('Shared fal prompt');
          expect(cacheKey).not.toContain('fal-tenant-a-secret');
          expect(cacheKey).not.toContain('fal-tenant-b-secret');
        }
      });

      it('should use a deterministic auth namespace across module reloads', async () => {
        async function getCacheKeyFromFreshModule() {
          vi.resetModules();
          const freshCache = await import('../../src/cache');
          const mockCache = {
            get: vi
              .fn()
              .mockResolvedValue(JSON.stringify('![cached](https://example.com/cached.png)')),
            set: vi.fn(),
          };
          vi.mocked(freshCache.isCacheEnabled).mockReturnValue(true);
          vi.mocked(freshCache.getCache).mockReturnValue(mockCache as any);
          const { FalImageGenerationProvider: FreshFalImageGenerationProvider } = await import(
            '../../src/providers/fal'
          );

          const freshProvider = new FreshFalImageGenerationProvider('fal-ai/flux/schnell', {
            config: { apiKey: 'fal-deterministic-secret' },
          });

          await freshProvider.callApi('Shared fal prompt');
          return mockCache.get.mock.calls[0][0] as string;
        }

        const firstCacheKey = await getCacheKeyFromFreshModule();
        const secondCacheKey = await getCacheKeyFromFreshModule();

        expect(firstCacheKey).toBe(secondCacheKey);
        expect(firstCacheKey).toMatch(
          /^fal:fal-ai\/flux\/schnell:[a-f0-9]{64}:[a-f0-9]{64}:[a-f0-9]{64}$/,
        );
        expect(firstCacheKey).not.toContain('fal-deterministic-secret');
        expect(firstCacheKey).not.toContain('Shared fal prompt');
      });

      it('should handle cache set errors gracefully', async () => {
        vi.mocked(isCacheEnabled).mockImplementation(function () {
          return true;
        });
        const mockCache = {
          get: vi.fn().mockResolvedValue(null),
          set: vi.fn().mockRejectedValue(new Error('Cache error')),
          wrap: vi.fn(),
          del: vi.fn(),
          clear: vi.fn(),
          stores: [
            {
              get: vi.fn(),
              set: vi.fn(),
            },
          ] as any,
          mget: vi.fn(),
          mset: vi.fn(),
          mdel: vi.fn(),
          reset: vi.fn(),
          ttl: vi.fn(),
          on: vi.fn(),
          removeAllListeners: vi.fn(),
        };
        vi.mocked(getCache).mockReturnValue(mockCache as any);

        const mockResponse = {
          data: {
            images: [{ url: 'https://example.com/image.png' }],
          },
          requestId: 'test-request-id',
        };
        mockSubscribe.mockResolvedValueOnce(mockResponse);

        const result = await provider.callApi('test prompt');

        expect(result).toEqual({
          cached: false,
          output: '![test prompt](https://example.com/image.png)',
        });
      });
    });

    describe('error handling', () => {
      it('should handle API errors', async () => {
        const mockError = new Error('API Error');
        mockSubscribe.mockRejectedValueOnce(mockError);

        await expect(provider.callApi('test prompt')).rejects.toThrow('API Error');
      });

      it('should throw error when image URL cannot be resolved', async () => {
        const mockResponse = {
          data: {
            // No images or image property
          },
          requestId: 'test-request-id',
        };
        mockSubscribe.mockResolvedValueOnce(mockResponse);

        await expect(
          provider.runInference({
            prompt: 'test prompt',
          }),
        ).rejects.toThrow('Failed to resolve image URL.');
      });
    });

    describe('client initialization', () => {
      it('should lazy load the fal client', async () => {
        vi.clearAllMocks();

        const newProvider = new FalImageGenerationProvider('fal-ai/flux/schnell', {
          config: { apiKey: 'test-api-key' },
        });

        const mockResponse = {
          data: {
            images: [{ url: 'https://example.com/image.png' }],
          },
          requestId: 'test-request-id',
        };
        mockSubscribe.mockResolvedValueOnce(mockResponse);

        await newProvider.callApi('test prompt');

        expect(mockSubscribe).toHaveBeenCalledWith('fal-ai/flux/schnell', {
          input: {
            prompt: 'test prompt',
          },
        });
        expect(mockConfig).toHaveBeenCalledWith({
          credentials: 'test-api-key',
        });
      });
    });
  });
});