TestGen
本项目基于assert语句对数据集进行扩充,并采用多条规则构建了焦点方法与测试方法之间具有准确匹配关系的高质量rust单元测试数据集。项目利用高质量数据集微调大语言模型,并结合静态分析方法构建思维链以引导模型推理,使模型生成rust单元测试的性能大幅提升。
使用说明
本项目流程分为以下几个步骤:
爬取并下载top-10k crates;
收集crates中的单元测试,并基于assert语句对单元测试数据集进行扩充;
基于匹配规则构建{焦点方法,测试方法}对数据集,匹配规则为
函数名匹配 且 (测试函数内调用或assert语句内调用);利用数据集对模型进行微调;
结合静态分析工具收集焦点方法上下文,构建思维链引导模型推理。焦点方法上下文包括:
- 调用焦点方法的方法;
- 焦点方法的父模块;
- 焦点方法内部调用的方法。
自动化验证生成的测试质量。
目录结构
crates_crawler: 爬取并下载top-10k crates的源代码;dataset_construction: 构建微调数据集的源代码;model_finetune: 模型微调的源代码;model_generation: 收集焦点方法上下文构建思维链以及模型推理的源代码;evaluation: 评估模型生成的单元测试质量的源代码。
Prerequisite
- 系统: Linux
- 环境
- python:
python==3.10 - python 依赖:
pip install -r requirements.txt - Rust 环境:
curl --proto '=https' --tlsv1.2 -sSf https://sh.rustup.rs | sh
- python:
- LLM
- 本地部署CodeLlama7B
运行方法
- 下载top-10k crates:进入
crates_crawler并执行python crates_crawler.py - 构建微调数据集:进入
dataset_construction/code并执行bash dataset_construction.sh; - 微调模型:进入
model_finetune/code并执行bash model_finetune.sh --model_path=<YourModelPath>; - 生成单元测试:进入
model_generation/code并执行bash model_generation.sh; - 评估:进入
evaluation/code并执行bash evaluation.sh。
评估结果
| 编译通过率 | 运行成功率 | |
|---|---|---|
| CodeLlama7B(original) | 19.88% | 6.45% |
| CodeLlama7B(fine-tuned) | 20.08% | 9.02% |
| CodeLlama7B(fine-tuned)+CoT | 25.41% | 13.42% |
TestGen VSCode插件
- VSCode扩展商店搜索
rusthelper并安装。
