ai-terminal-assistant/tests/unit/tools/web/web_extract.test.ts

import { describe, it, expect, beforeEach, vi } from 'vitest';

// Mock tavily
const mockExtract = vi.fn();
vi.mock('@tavily/core', () => ({
  tavily: vi.fn(() => ({
    extract: mockExtract,
  })),
}));

// Mock config
vi.mock('../../../../src/utils/config.js', () => ({
  getConfig: vi.fn(() => ({
    tavilyApiKey: 'test-api-key',
  })),
}));

// Mock permission manager
vi.mock('../../../../src/permission/index.js', () => ({
  getPermissionManager: vi.fn(() => ({
    checkWebPermission: vi.fn().mockResolvedValue({
      allowed: true,
      action: 'allow',
    }),
  })),
}));

// Mock loadDescription
vi.mock('../../../../src/tools/load_description.js', () => ({
  loadDescription: vi.fn(() => '从网页URL提取内容'),
}));

import { webExtractTool } from '../../../../src/tools/web/web_extract.js';
import { getPermissionManager } from '../../../../src/permission/index.js';
import { getConfig } from '../../../../src/utils/config.js';

describe('webExtractTool - 网页内容提取工具', () => {
  beforeEach(() => {
    vi.clearAllMocks();
    mockExtract.mockResolvedValue({
      results: [
        {
          url: 'https://example.com',
          rawContent: '# Hello World\n\nThis is example content.',
          images: ['https://example.com/img1.png'],
        },
      ],
      failedResults: [],
      responseTime: 1.5,
    });
  });

  describe('工具定义', () => {
    it('有正确的名称', () => {
      expect(webExtractTool.name).toBe('web_extract');
    });

    it('有正确的元数据', () => {
      expect(webExtractTool.metadata.category).toBe('web');
      expect(webExtractTool.metadata.keywords).toContain('extract');
      expect(webExtractTool.metadata.keywords).toContain('url');
      expect(webExtractTool.metadata.keywords).toContain('scrape');
    });

    it('定义了必需的 urls 参数', () => {
      expect(webExtractTool.parameters.urls.required).toBe(true);
    });

    it('定义了可选参数', () => {
      expect(webExtractTool.parameters.extract_depth.required).toBe(false);
      expect(webExtractTool.parameters.format.required).toBe(false);
      expect(webExtractTool.parameters.include_images.required).toBe(false);
    });
  });

  describe('execute - 执行', () => {
    it('成功提取单个 URL 内容', async () => {
      const result = await webExtractTool.execute({
        urls: ['https://example.com'],
      });

      expect(result.success).toBe(true);
      expect(result.output).toContain('网页内容提取');
      expect(result.output).toContain('example.com');
      expect(result.output).toContain('Hello World');
    });

    it('支持字符串格式的单个 URL', async () => {
      const result = await webExtractTool.execute({
        urls: 'https://example.com',
      });

      expect(result.success).toBe(true);
      expect(mockExtract).toHaveBeenCalledWith(
        ['https://example.com'],
        expect.any(Object)
      );
    });

    it('支持多个 URL', async () => {
      mockExtract.mockResolvedValue({
        results: [
          { url: 'https://example1.com', rawContent: 'Content 1' },
          { url: 'https://example2.com', rawContent: 'Content 2' },
        ],
        failedResults: [],
      });

      const result = await webExtractTool.execute({
        urls: ['https://example1.com', 'https://example2.com'],
      });

      expect(result.success).toBe(true);
      expect(result.output).toContain('example1.com');
      expect(result.output).toContain('example2.com');
    });

    it('限制最多 20 个 URL', async () => {
      const urls = Array.from({ length: 25 }, (_, i) => `https://example${i}.com`);

      await webExtractTool.execute({ urls });

      expect(mockExtract).toHaveBeenCalledWith(
        expect.any(Array),
        expect.any(Object)
      );
      const calledUrls = mockExtract.mock.calls[0][0];
      expect(calledUrls.length).toBe(20);
    });

    it('使用 advanced 提取深度', async () => {
      await webExtractTool.execute({
        urls: ['https://example.com'],
        extract_depth: 'advanced',
      });

      expect(mockExtract).toHaveBeenCalledWith(
        expect.any(Array),
        expect.objectContaining({ extractDepth: 'advanced' })
      );
    });

    it('包含图片列表', async () => {
      const result = await webExtractTool.execute({
        urls: ['https://example.com'],
        include_images: true,
      });

      expect(result.success).toBe(true);
      expect(result.output).toContain('图片');
      expect(result.output).toContain('img1.png');
    });

    it('显示失败的 URL', async () => {
      mockExtract.mockResolvedValue({
        results: [],
        failedResults: [
          { url: 'https://failed.com', error: '404 Not Found' },
        ],
      });

      const result = await webExtractTool.execute({
        urls: ['https://failed.com'],
      });

      expect(result.success).toBe(true);
      expect(result.output).toContain('提取失败');
      expect(result.output).toContain('failed.com');
      expect(result.output).toContain('404');
    });

    it('显示响应时间', async () => {
      const result = await webExtractTool.execute({
        urls: ['https://example.com'],
      });

      expect(result.success).toBe(true);
      expect(result.output).toContain('提取耗时');
    });

    it('截断过长的内容', async () => {
      mockExtract.mockResolvedValue({
        results: [
          {
            url: 'https://example.com',
            rawContent: 'a'.repeat(6000), // 超过 5000 字符
          },
        ],
        failedResults: [],
      });

      const result = await webExtractTool.execute({
        urls: ['https://example.com'],
      });

      expect(result.success).toBe(true);
      expect(result.output).toContain('内容已截断');
    });

    it('无 API Key 返回错误', async () => {
      vi.mocked(getConfig).mockReturnValue({} as any);
      const originalEnv = process.env.TAVILY_API_KEY;
      delete process.env.TAVILY_API_KEY;

      const result = await webExtractTool.execute({
        urls: ['https://example.com'],
      });

      expect(result.success).toBe(false);
      expect(result.error).toContain('未配置 Tavily API Key');

      process.env.TAVILY_API_KEY = originalEnv;
    });

    it('权限被拒绝时返回错误', async () => {
      vi.mocked(getPermissionManager).mockReturnValue({
        checkWebPermission: vi.fn().mockResolvedValue({
          allowed: false,
          action: 'deny',
          reason: '提取不被允许',
        }),
      } as any);

      const result = await webExtractTool.execute({
        urls: ['https://example.com'],
      });

      expect(result.success).toBe(false);
      expect(result.error).toContain('权限被拒绝');
    });

    it('需要确认时返回提示', async () => {
      vi.mocked(getPermissionManager).mockReturnValue({
        checkWebPermission: vi.fn().mockResolvedValue({
          allowed: false,
          action: 'ask',
          needsConfirmation: true,
        }),
      } as any);

      const result = await webExtractTool.execute({
        urls: ['https://example.com'],
      });

      expect(result.success).toBe(false);
      expect(result.error).toContain('需要用户确认');
    });

    it('提取失败返回错误', async () => {
      vi.mocked(getPermissionManager).mockReturnValue({
        checkWebPermission: vi.fn().mockResolvedValue({ allowed: true }),
      } as any);
      mockExtract.mockRejectedValue(new Error('API 调用失败'));

      const result = await webExtractTool.execute({
        urls: ['https://example.com'],
      });

      expect(result.success).toBe(false);
      expect(result.error).toContain('提取失败');
      expect(result.error).toContain('API 调用失败');
    });
  });
});