ai-terminal-assistant/packages/core/tests/repomap/repomap.test.ts

/**
 * RepoMap 测试
 */

import { describe, it, expect, beforeAll } from 'vitest';
import * as path from 'path';
import * as fs from 'fs/promises';
import * as os from 'os';
import { createRepoMap } from '../../src/repomap/repomap.js';
import { TagExtractor } from '../../src/repomap/tags/extractor.js';
import { Graph, pagerank, distributeRanksToDefinitions } from '../../src/repomap/ranking/index.js';
import { DiskCache } from '../../src/repomap/cache/disk-cache.js';

describe('Graph', () => {
  it('should add edges and track nodes', () => {
    const graph = new Graph();

    graph.addEdge({ from: 'a.ts', to: 'b.ts', weight: 1, ident: 'foo' });
    graph.addEdge({ from: 'a.ts', to: 'c.ts', weight: 2, ident: 'bar' });
    graph.addEdge({ from: 'b.ts', to: 'c.ts', weight: 1, ident: 'baz' });

    const nodes = graph.getNodes();
    expect(nodes).toContain('a.ts');
    expect(nodes).toContain('b.ts');
    expect(nodes).toContain('c.ts');
    expect(nodes.length).toBe(3);
  });

  it('should track in and out edges correctly', () => {
    const graph = new Graph();

    graph.addEdge({ from: 'a.ts', to: 'b.ts', weight: 1, ident: 'foo' });
    graph.addEdge({ from: 'a.ts', to: 'c.ts', weight: 2, ident: 'bar' });

    const outEdges = graph.getOutEdges('a.ts');
    expect(outEdges.length).toBe(2);

    const inEdges = graph.getInEdges('b.ts');
    expect(inEdges.length).toBe(1);
    expect(inEdges[0].from).toBe('a.ts');
  });

  it('should calculate out degree correctly', () => {
    const graph = new Graph();

    graph.addEdge({ from: 'a.ts', to: 'b.ts', weight: 1, ident: 'foo' });
    graph.addEdge({ from: 'a.ts', to: 'c.ts', weight: 2, ident: 'bar' });

    expect(graph.getOutDegree('a.ts')).toBe(3); // sum of weights
    expect(graph.getOutDegree('b.ts')).toBe(0);
  });
});

describe('PageRank', () => {
  it('should compute ranks for a simple graph', () => {
    const graph = new Graph();

    // 创建一个简单的图: a -> b -> c
    graph.addEdge({ from: 'a.ts', to: 'b.ts', weight: 1, ident: 'foo' });
    graph.addEdge({ from: 'b.ts', to: 'c.ts', weight: 1, ident: 'bar' });

    const ranks = pagerank(graph);

    expect(ranks.size).toBe(3);
    expect(ranks.get('a.ts')).toBeGreaterThan(0);
    expect(ranks.get('b.ts')).toBeGreaterThan(0);
    expect(ranks.get('c.ts')).toBeGreaterThan(0);
  });

  it('should respect personalization vector', () => {
    const graph = new Graph();

    graph.addEdge({ from: 'a.ts', to: 'b.ts', weight: 1, ident: 'foo' });
    graph.addEdge({ from: 'b.ts', to: 'c.ts', weight: 1, ident: 'bar' });
    graph.addEdge({ from: 'c.ts', to: 'a.ts', weight: 1, ident: 'baz' });

    // 高度偏好 a.ts
    const personalization = new Map<string, number>();
    personalization.set('a.ts', 100);
    personalization.set('b.ts', 1);
    personalization.set('c.ts', 1);

    const ranks = pagerank(graph, { personalization });

    // a.ts 应该有较高的 rank
    expect(ranks.get('a.ts')!).toBeGreaterThan(ranks.get('b.ts')!);
  });

  it('should distribute ranks to definitions', () => {
    const graph = new Graph();

    graph.addEdge({ from: 'a.ts', to: 'b.ts', weight: 1, ident: 'foo' });
    graph.addEdge({ from: 'a.ts', to: 'b.ts', weight: 2, ident: 'bar' });

    const nodeRanks = pagerank(graph);
    const defRanks = distributeRanksToDefinitions(graph, nodeRanks);

    expect(defRanks.has('b.ts:foo')).toBe(true);
    expect(defRanks.has('b.ts:bar')).toBe(true);
  });
});

describe('DiskCache', () => {
  let cacheDir: string;
  let cache: DiskCache<{ value: number }>;

  beforeAll(async () => {
    cacheDir = path.join(os.tmpdir(), `repomap-test-${Date.now()}`);
    cache = new DiskCache(cacheDir);
  });

  it('should set and get values', async () => {
    await cache.set('key1', { value: 42 });
    const result = await cache.get('key1');

    expect(result).not.toBeNull();
    expect(result?.value).toBe(42);
  });

  it('should return null for missing keys', async () => {
    const result = await cache.get('nonexistent');
    expect(result).toBeNull();
  });

  it('should delete values', async () => {
    await cache.set('key2', { value: 100 });
    await cache.delete('key2');

    const result = await cache.get('key2');
    expect(result).toBeNull();
  });

  it('should check existence with has', async () => {
    await cache.set('key3', { value: 200 });

    expect(await cache.has('key3')).toBe(true);
    expect(await cache.has('nonexistent2')).toBe(false);
  });
});

describe('TagExtractor', () => {
  let tempDir: string;

  beforeAll(async () => {
    tempDir = path.join(os.tmpdir(), `repomap-extractor-test-${Date.now()}`);
    await fs.mkdir(tempDir, { recursive: true });
  });

  it('should extract tags from TypeScript code using regex fallback', async () => {
    const testFile = path.join(tempDir, 'test.ts');
    const code = `
export function greet(name: string): string {
  return \`Hello, \${name}!\`;
}

export class Greeter {
  private name: string;

  constructor(name: string) {
    this.name = name;
  }

  greet(): string {
    return \`Hello, \${this.name}!\`;
  }
}

export interface IGreeter {
  greet(): string;
}
`;

    await fs.writeFile(testFile, code);

    const extractor = new TagExtractor();
    const tags = await extractor.getTags(testFile, 'test.ts');

    // 检查是否提取到了一些 tags
    expect(tags.length).toBeGreaterThan(0);

    // 检查是否有定义
    const defs = tags.filter((t) => t.kind === 'def');
    expect(defs.length).toBeGreaterThan(0);

    // 检查是否包含函数和类
    const defNames = defs.map((t) => t.name);
    expect(defNames).toContain('greet');
    expect(defNames).toContain('Greeter');
  });

  it('should extract tags from Python code using regex fallback', async () => {
    const testFile = path.join(tempDir, 'test.py');
    const code = `
def greet(name: str) -> str:
    return f"Hello, {name}!"

class Greeter:
    def __init__(self, name: str):
        self.name = name

    def greet(self) -> str:
        return f"Hello, {self.name}!"

async def async_greet(name: str) -> str:
    return f"Hello async, {name}!"
`;

    await fs.writeFile(testFile, code);

    const extractor = new TagExtractor();
    const tags = await extractor.getTags(testFile, 'test.py');

    // 检查是否提取到了一些 tags
    expect(tags.length).toBeGreaterThan(0);

    // 检查是否有定义
    const defs = tags.filter((t) => t.kind === 'def');
    expect(defs.length).toBeGreaterThan(0);

    // 检查是否包含函数和类
    const defNames = defs.map((t) => t.name);
    expect(defNames).toContain('greet');
    expect(defNames).toContain('Greeter');
    expect(defNames).toContain('async_greet');
  });
});

describe('RepoMap', () => {
  let tempDir: string;

  beforeAll(async () => {
    tempDir = path.join(os.tmpdir(), `repomap-test-${Date.now()}`);
    await fs.mkdir(tempDir, { recursive: true });

    // 创建测试文件
    await fs.writeFile(
      path.join(tempDir, 'utils.ts'),
      `
export function add(a: number, b: number): number {
  return a + b;
}

export function multiply(a: number, b: number): number {
  return a * b;
}
`
    );

    await fs.writeFile(
      path.join(tempDir, 'calculator.ts'),
      `
import { add, multiply } from './utils';

export class Calculator {
  add(a: number, b: number): number {
    return add(a, b);
  }

  multiply(a: number, b: number): number {
    return multiply(a, b);
  }
}
`
    );

    await fs.writeFile(
      path.join(tempDir, 'main.ts'),
      `
import { Calculator } from './calculator';

const calc = new Calculator();
console.log(calc.add(1, 2));
console.log(calc.multiply(3, 4));
`
    );
  });

  it('should create a repo map', async () => {
    const repoMap = createRepoMap(tempDir, {
      mapTokens: 2048,
    });

    const files = [
      path.join(tempDir, 'utils.ts'),
      path.join(tempDir, 'calculator.ts'),
      path.join(tempDir, 'main.ts'),
    ];

    const map = await repoMap.getRepoMap([], files, new Set(), new Set());

    // 应该生成了一些内容
    expect(map.length).toBeGreaterThan(0);
  });

  it('should boost mentioned identifiers', async () => {
    const repoMap = createRepoMap(tempDir, {
      mapTokens: 2048,
    });

    const files = [
      path.join(tempDir, 'utils.ts'),
      path.join(tempDir, 'calculator.ts'),
      path.join(tempDir, 'main.ts'),
    ];

    const mentionedIdents = new Set(['Calculator']);
    const map = await repoMap.getRepoMap([], files, new Set(), mentionedIdents);

    // 应该包含 Calculator 相关内容
    expect(map.length).toBeGreaterThan(0);
  });

  it('should exclude chat files from output', async () => {
    const repoMap = createRepoMap(tempDir, {
      mapTokens: 2048,
    });

    const chatFiles = [path.join(tempDir, 'utils.ts')];
    const otherFiles = [
      path.join(tempDir, 'calculator.ts'),
      path.join(tempDir, 'main.ts'),
    ];

    const map = await repoMap.getRepoMap(chatFiles, otherFiles, new Set(), new Set());

    // utils.ts 不应该出现在输出中
    expect(map).not.toContain('utils.ts:');
  });
});