chore: import upstream snapshot with attribution

test/integrations/huggingfaceDatasets.test.ts

@@ -0,0 +1,586 @@
+import { afterEach, beforeEach, describe, expect, it, vi } from 'vitest';
+import { fetchWithCache } from '../../src/cache';
+import { getEnvString } from '../../src/envars';
+import {
+  fetchHuggingFaceDataset,
+  parseDatasetPath,
+} from '../../src/integrations/huggingfaceDatasets';
+
+vi.mock('../../src/cache', () => ({
+  fetchWithCache: vi.fn(),
+}));
+
+vi.mock('../../src/util/fetch/index.ts', () => ({
+  fetchWithProxy: vi.fn(),
+}));
+
+vi.mock('../../src/envars', () => ({
+  getEnvString: vi.fn().mockReturnValue(''),
+  isCI: vi.fn().mockReturnValue(false),
+}));
+
+describe('huggingfaceDatasets', () => {
+  beforeEach(() => {
+    vi.mocked(getEnvString).mockReturnValue('');
+  });
+
+  afterEach(() => {
+    vi.resetAllMocks();
+  });
+
+  describe('parseDatasetPath', () => {
+    it('should parse path with default parameters', () => {
+      const result = parseDatasetPath('huggingface://datasets/owner/repo');
+      expect(result).toEqual({
+        owner: 'owner',
+        repo: 'repo',
+        queryParams: expect.any(URLSearchParams),
+      });
+      expect(result.queryParams.get('split')).toBe('test');
+      expect(result.queryParams.get('config')).toBe('default');
+    });
+
+    it('should parse path with custom query parameters', () => {
+      const result = parseDatasetPath(
+        'huggingface://datasets/owner/repo?split=train&config=custom&limit=10',
+      );
+      expect(result).toEqual({
+        owner: 'owner',
+        repo: 'repo',
+        queryParams: expect.any(URLSearchParams),
+      });
+      expect(result.queryParams.get('split')).toBe('train');
+      expect(result.queryParams.get('config')).toBe('custom');
+      expect(result.queryParams.get('limit')).toBe('10');
+    });
+
+    it('should override default parameters with user parameters', () => {
+      const result = parseDatasetPath('huggingface://datasets/owner/repo?split=validation');
+      expect(result.queryParams.get('split')).toBe('validation');
+      expect(result.queryParams.get('config')).toBe('default');
+    });
+  });
+
+  it('should fetch and parse dataset with default parameters', async () => {
+    vi.mocked(fetchWithCache).mockResolvedValueOnce({
+      data: {
+        num_rows_total: 2,
+        features: [
+          { name: 'act', type: { dtype: 'string', _type: 'Value' } },
+          { name: 'prompt', type: { dtype: 'string', _type: 'Value' } },
+        ],
+        rows: [
+          { row: { act: 'Linux Terminal', prompt: 'List all files' } },
+          { row: { act: 'Math Tutor', prompt: 'Solve 2+2' } },
+        ],
+      },
+      cached: false,
+      status: 200,
+      statusText: 'OK',
+    } as any);
+
+    const tests = await fetchHuggingFaceDataset('huggingface://datasets/test/dataset');
+
+    expect(vi.mocked(fetchWithCache)).toHaveBeenCalledWith(
+      'https://datasets-server.huggingface.co/rows?dataset=test%2Fdataset&split=test&config=default&offset=0&length=100',
+      expect.objectContaining({
+        headers: {},
+      }),
+    );
+
+    expect(tests).toHaveLength(2);
+    expect(tests[0].vars).toEqual({
+      act: 'Linux Terminal',
+      prompt: 'List all files',
+    });
+    expect(tests[1].vars).toEqual({
+      act: 'Math Tutor',
+      prompt: 'Solve 2+2',
+    });
+
+    // Check that disableVarExpansion is set for all test cases
+    tests.forEach((test) => {
+      expect(test.options).toEqual({
+        disableVarExpansion: true,
+      });
+    });
+  });
+
+  it('should include auth token when HF_TOKEN is set', async () => {
+    vi.mocked(getEnvString).mockImplementation((key) => {
+      if (key === 'HF_TOKEN') {
+        return 'test-token';
+      }
+      return '';
+    });
+
+    vi.mocked(fetchWithCache).mockResolvedValueOnce({
+      data: {
+        num_rows_total: 1,
+        features: [
+          { name: 'question', type: { dtype: 'string', _type: 'Value' } },
+          { name: 'answer', type: { dtype: 'string', _type: 'Value' } },
+        ],
+        rows: [{ row: { question: 'What is 2+2?', answer: '4' } }],
+      },
+      cached: false,
+      status: 200,
+      statusText: 'OK',
+    } as any);
+
+    await fetchHuggingFaceDataset('huggingface://datasets/test/dataset');
+
+    expect(vi.mocked(fetchWithCache)).toHaveBeenCalledWith(
+      'https://datasets-server.huggingface.co/rows?dataset=test%2Fdataset&split=test&config=default&offset=0&length=100',
+      expect.objectContaining({
+        headers: {
+          Authorization: 'Bearer test-token',
+        },
+      }),
+    );
+  });
+
+  it('should fall back to HF_API_TOKEN when HF_TOKEN is empty', async () => {
+    vi.mocked(getEnvString).mockImplementation((key) => {
+      if (key === 'HF_TOKEN') {
+        return '';
+      }
+      if (key === 'HF_API_TOKEN') {
+        return 'api-token';
+      }
+      return '';
+    });
+
+    vi.mocked(fetchWithCache).mockResolvedValueOnce({
+      data: {
+        num_rows_total: 1,
+        features: [{ name: 'text', type: { dtype: 'string', _type: 'Value' } }],
+        rows: [{ row: { text: 'test' } }],
+      },
+      cached: false,
+      status: 200,
+      statusText: 'OK',
+    } as any);
+
+    await fetchHuggingFaceDataset('huggingface://datasets/test/dataset', 1);
+
+    expect(vi.mocked(fetchWithCache)).toHaveBeenCalledWith(
+      expect.any(String),
+      expect.objectContaining({
+        headers: {
+          Authorization: 'Bearer api-token',
+        },
+      }),
+    );
+  });
+
+  it('should fall back to HUGGING_FACE_HUB_TOKEN when other tokens are empty', async () => {
+    vi.mocked(getEnvString).mockImplementation((key) => {
+      if (key === 'HF_TOKEN') {
+        return '';
+      }
+      if (key === 'HF_API_TOKEN') {
+        return '';
+      }
+      if (key === 'HUGGING_FACE_HUB_TOKEN') {
+        return 'hub-token';
+      }
+      return '';
+    });
+
+    vi.mocked(fetchWithCache).mockResolvedValueOnce({
+      data: {
+        num_rows_total: 1,
+        features: [{ name: 'text', type: { dtype: 'string', _type: 'Value' } }],
+        rows: [{ row: { text: 'test' } }],
+      },
+      cached: false,
+      status: 200,
+      statusText: 'OK',
+    } as any);
+
+    await fetchHuggingFaceDataset('huggingface://datasets/test/dataset', 1);
+
+    expect(vi.mocked(fetchWithCache)).toHaveBeenCalledWith(
+      expect.any(String),
+      expect.objectContaining({
+        headers: {
+          Authorization: 'Bearer hub-token',
+        },
+      }),
+    );
+  });
+
+  it('should handle custom query parameters', async () => {
+    vi.mocked(fetchWithCache).mockResolvedValueOnce({
+      data: {
+        num_rows_total: 1,
+        features: [
+          { name: 'question', type: { dtype: 'string', _type: 'Value' } },
+          { name: 'answer', type: { dtype: 'string', _type: 'Value' } },
+        ],
+        rows: [{ row: { question: 'What is 2+2?', answer: '4' } }],
+      },
+      cached: false,
+      status: 200,
+      statusText: 'OK',
+    } as any);
+
+    await fetchHuggingFaceDataset('huggingface://datasets/test/dataset?split=train&config=custom');
+
+    expect(vi.mocked(fetchWithCache)).toHaveBeenCalledWith(
+      'https://datasets-server.huggingface.co/rows?dataset=test%2Fdataset&split=train&config=custom&offset=0&length=100',
+      expect.objectContaining({
+        headers: {},
+      }),
+    );
+  });
+
+  it('should handle pagination', async () => {
+    vi.mocked(fetchWithCache).mockResolvedValueOnce({
+      data: {
+        num_rows_total: 3,
+        features: [{ name: 'text', type: { dtype: 'string', _type: 'Value' } }],
+        rows: [{ row: { text: 'First' } }, { row: { text: 'Second' } }],
+      },
+      cached: false,
+      status: 200,
+      statusText: 'OK',
+    } as any);
+
+    vi.mocked(fetchWithCache).mockResolvedValueOnce({
+      data: {
+        num_rows_total: 3,
+        features: [{ name: 'text', type: { dtype: 'string', _type: 'Value' } }],
+        rows: [{ row: { text: 'Third' } }],
+      },
+      cached: false,
+      status: 200,
+      statusText: 'OK',
+    } as any);
+
+    const tests = await fetchHuggingFaceDataset('huggingface://datasets/test/dataset');
+
+    expect(vi.mocked(fetchWithCache)).toHaveBeenCalledTimes(2);
+    expect(vi.mocked(fetchWithCache)).toHaveBeenNthCalledWith(
+      1,
+      'https://datasets-server.huggingface.co/rows?dataset=test%2Fdataset&split=test&config=default&offset=0&length=100',
+      expect.objectContaining({
+        headers: {},
+      }),
+    );
+    // Note: Second call might have different length due to concurrent fetching and remaining calculation
+    expect(vi.mocked(fetchWithCache)).toHaveBeenNthCalledWith(
+      2,
+      expect.stringContaining('dataset=test%2Fdataset&split=test&config=default&offset=2'),
+      expect.objectContaining({
+        headers: {},
+      }),
+    );
+
+    expect(tests).toHaveLength(3);
+    expect(tests.map((t) => t.vars?.text)).toEqual(['First', 'Second', 'Third']);
+
+    // Check that disableVarExpansion is set for all test cases
+    tests.forEach((test) => {
+      expect(test.options).toEqual({
+        disableVarExpansion: true,
+      });
+    });
+  });
+
+  it('should reject an empty non-final page instead of retrying the same offset', async () => {
+    vi.mocked(fetchWithCache)
+      .mockResolvedValueOnce({
+        data: {
+          num_rows_total: 2,
+          features: [{ name: 'text', type: { dtype: 'string', _type: 'Value' } }],
+          rows: [],
+        },
+        cached: false,
+        status: 200,
+        statusText: 'OK',
+      } as any)
+      .mockRejectedValueOnce(new Error('Paginator retried the empty page'));
+
+    await expect(
+      fetchHuggingFaceDataset('huggingface://datasets/test/dataset', 200),
+    ).rejects.toThrow(
+      '[HF Dataset] Received an empty page at offset 0 before reaching 2 total rows',
+    );
+
+    expect(vi.mocked(fetchWithCache)).toHaveBeenCalledTimes(1);
+  });
+
+  it('should handle API errors by throwing', async () => {
+    vi.mocked(fetchWithCache).mockResolvedValueOnce({
+      data: null,
+      cached: false,
+      status: 404,
+      statusText: 'Not Found',
+    } as any);
+
+    await expect(
+      fetchHuggingFaceDataset('huggingface://datasets/nonexistent/dataset'),
+    ).rejects.toThrow('[HF Dataset] Failed to fetch dataset: Not Found');
+  });
+
+  it('should short-circuit and return [] when limit is 0', async () => {
+    const tests = await fetchHuggingFaceDataset('huggingface://datasets/test/dataset', 0);
+    expect(vi.mocked(fetchWithCache)).not.toHaveBeenCalled();
+    expect(tests).toEqual([]);
+  });
+
+  it('should respect user-specified limit parameter (single request optimization)', async () => {
+    vi.mocked(fetchWithCache).mockResolvedValueOnce({
+      data: {
+        num_rows_total: 5,
+        features: [{ name: 'text', type: { dtype: 'string', _type: 'Value' } }],
+        rows: [{ row: { text: 'First' } }, { row: { text: 'Second' } }],
+      },
+      cached: false,
+      status: 200,
+      statusText: 'OK',
+    } as any);
+
+    const tests = await fetchHuggingFaceDataset('huggingface://datasets/test/dataset?limit=2');
+
+    expect(vi.mocked(fetchWithCache)).toHaveBeenCalledWith(
+      'https://datasets-server.huggingface.co/rows?dataset=test%2Fdataset&split=test&config=default&limit=2&offset=0&length=2',
+      expect.objectContaining({
+        headers: {},
+      }),
+    );
+
+    expect(tests).toHaveLength(2);
+    expect(tests.map((t) => t.vars?.text)).toEqual(['First', 'Second']);
+
+    // Check that disableVarExpansion is set for all test cases
+    tests.forEach((test) => {
+      expect(test.options).toEqual({
+        disableVarExpansion: true,
+      });
+    });
+  });
+
+  it('should handle limit larger than page size', async () => {
+    vi.mocked(fetchWithCache).mockResolvedValueOnce({
+      data: {
+        num_rows_total: 150,
+        features: [{ name: 'text', type: { dtype: 'string', _type: 'Value' } }],
+        rows: Array(100)
+          .fill(null)
+          .map((_, i) => ({ row: { text: `Item ${i + 1}` } })),
+      },
+      cached: false,
+      status: 200,
+      statusText: 'OK',
+    } as any);
+
+    vi.mocked(fetchWithCache).mockResolvedValueOnce({
+      data: {
+        num_rows_total: 150,
+        features: [{ name: 'text', type: { dtype: 'string', _type: 'Value' } }],
+        rows: Array(20)
+          .fill(null)
+          .map((_, i) => ({ row: { text: `Item ${i + 101}` } })),
+      },
+      cached: false,
+      status: 200,
+      statusText: 'OK',
+    } as any);
+
+    const tests = await fetchHuggingFaceDataset('huggingface://datasets/test/dataset?limit=120');
+
+    expect(vi.mocked(fetchWithCache)).toHaveBeenCalledTimes(2);
+    expect(vi.mocked(fetchWithCache)).toHaveBeenNthCalledWith(
+      1,
+      'https://datasets-server.huggingface.co/rows?dataset=test%2Fdataset&split=test&config=default&limit=120&offset=0&length=100',
+      expect.objectContaining({
+        headers: {},
+      }),
+    );
+    expect(vi.mocked(fetchWithCache)).toHaveBeenNthCalledWith(
+      2,
+      'https://datasets-server.huggingface.co/rows?dataset=test%2Fdataset&split=test&config=default&limit=120&offset=100&length=20',
+      expect.objectContaining({
+        headers: {},
+      }),
+    );
+
+    expect(tests).toHaveLength(120);
+    expect(tests[119].vars?.text).toBe('Item 120');
+
+    // Check that disableVarExpansion is set for all test cases
+    tests.forEach((test) => {
+      expect(test.options).toEqual({
+        disableVarExpansion: true,
+      });
+    });
+  });
+
+  describe('performance optimizations', () => {
+    it('should use single request optimization for small limits', async () => {
+      vi.mocked(fetchWithCache).mockResolvedValueOnce({
+        data: {
+          num_rows_total: 1000,
+          features: [{ name: 'text', type: { dtype: 'string', _type: 'Value' } }],
+          rows: Array(50)
+            .fill(null)
+            .map((_, i) => ({ row: { text: `Item ${i + 1}` } })),
+        },
+        cached: true,
+        status: 200,
+        statusText: 'OK',
+      } as any);
+
+      const tests = await fetchHuggingFaceDataset('huggingface://datasets/test/dataset', 50);
+
+      // Should only make one request for limits <= 100
+      expect(vi.mocked(fetchWithCache)).toHaveBeenCalledTimes(1);
+      expect(tests).toHaveLength(50);
+    });
+
+    it('should throw error on page fetch failure', async () => {
+      // First page succeeds
+      vi.mocked(fetchWithCache).mockResolvedValueOnce({
+        data: {
+          num_rows_total: 300,
+          features: [{ name: 'text', type: { dtype: 'string', _type: 'Value' } }],
+          rows: Array(100)
+            .fill(null)
+            .map((_, i) => ({ row: { text: `Item ${i + 1}` } })),
+        },
+        cached: false,
+        status: 200,
+        statusText: 'OK',
+      } as any);
+
+      // Second page fails
+      vi.mocked(fetchWithCache).mockResolvedValueOnce({
+        data: null,
+        cached: false,
+        status: 500,
+        statusText: 'Internal Server Error',
+      } as any);
+
+      // Should throw error instead of returning partial results
+      await expect(
+        fetchHuggingFaceDataset('huggingface://datasets/test/dataset', 200),
+      ).rejects.toThrow('[HF Dataset] Failed to fetch dataset: Internal Server Error');
+    });
+
+    it('should include rows from concurrent fetches without duplicates', async () => {
+      const totalRows = 400;
+      const rowPrefix = 'x'.repeat(300);
+
+      vi.mocked(fetchWithCache).mockImplementation(async (url) => {
+        const searchParams = new URL(String(url)).searchParams;
+        const offset = Number.parseInt(searchParams.get('offset') ?? '0', 10);
+        const length = Number.parseInt(searchParams.get('length') ?? '100', 10);
+
+        return {
+          data: {
+            num_rows_total: totalRows,
+            features: [{ name: 'text', type: { dtype: 'string', _type: 'Value' } }],
+            rows: Array.from({ length }, (_, i) => ({
+              row: { text: `${rowPrefix}${offset + i + 1}` },
+            })),
+          },
+          cached: false,
+          status: 200,
+          statusText: 'OK',
+        } as any;
+      });
+
+      const tests = await fetchHuggingFaceDataset('huggingface://datasets/test/dataset');
+      const texts = tests.map((test) => test.vars?.text);
+      const calledOffsets = vi
+        .mocked(fetchWithCache)
+        .mock.calls.map(([url]) => new URL(String(url)).searchParams.get('offset'));
+
+      expect(calledOffsets).toContain('100');
+      expect(tests).toHaveLength(totalRows);
+      expect(texts).toContain(`${rowPrefix}150`);
+      expect(new Set(texts).size).toBe(totalRows);
+    });
+
+    it('should adapt page size based on row size', async () => {
+      // Mock a dataset with large rows (>2KB each)
+      const largeRow = { text: 'x'.repeat(3000) }; // ~3KB row
+
+      vi.mocked(fetchWithCache).mockResolvedValueOnce({
+        data: {
+          num_rows_total: 200,
+          features: [{ name: 'text', type: { dtype: 'string', _type: 'Value' } }],
+          rows: Array(100)
+            .fill(null)
+            .map(() => ({ row: largeRow })),
+        },
+        cached: false,
+        status: 200,
+        statusText: 'OK',
+      } as any);
+
+      // Mock all subsequent potential concurrent requests to avoid undefined errors
+      vi.mocked(fetchWithCache).mockResolvedValue({
+        data: {
+          num_rows_total: 200,
+          features: [{ name: 'text', type: { dtype: 'string', _type: 'Value' } }],
+          rows: Array(25)
+            .fill(null)
+            .map((_, i) => ({ row: { text: `Item ${i + 101}` } })),
+        },
+        cached: false,
+        status: 200,
+        statusText: 'OK',
+      } as any);
+
+      const tests = await fetchHuggingFaceDataset('huggingface://datasets/test/dataset', 125);
+
+      // Should have made at least one request
+      expect(vi.mocked(fetchWithCache)).toHaveBeenCalled();
+
+      // First call should be normal page size
+      expect(vi.mocked(fetchWithCache)).toHaveBeenNthCalledWith(
+        1,
+        expect.stringContaining('length=100'),
+        expect.anything(),
+      );
+
+      expect(tests.length).toBeGreaterThan(0);
+      expect(tests[0].vars?.text).toBe(largeRow.text);
+    });
+
+    it('should handle authentication tokens correctly', async () => {
+      vi.mocked(getEnvString).mockImplementation((key) => {
+        if (key === 'HF_TOKEN') {
+          return 'test-token-123';
+        }
+        return '';
+      });
+
+      vi.mocked(fetchWithCache).mockResolvedValueOnce({
+        data: {
+          num_rows_total: 10,
+          features: [{ name: 'text', type: { dtype: 'string', _type: 'Value' } }],
+          rows: [{ row: { text: 'Test' } }],
+        },
+        cached: false,
+        status: 200,
+        statusText: 'OK',
+      } as any);
+
+      await fetchHuggingFaceDataset('huggingface://datasets/test/dataset', 5);
+
+      expect(vi.mocked(fetchWithCache)).toHaveBeenCalledWith(
+        expect.any(String),
+        expect.objectContaining({
+          headers: {
+            Authorization: 'Bearer test-token-123',
+          },
+        }),
+      );
+    });
+  });
+});