Improving Deep Assertion Generation via Fine-Tuning Retrieval-Augmented Pre-trained Language Models

2025-02-22

AI Summary

Existing assertion generation methods rely heavily on lexical matching and suffer from limited training data, leading to inadequate semantic understanding and low accuracy. To address these limitations, we propose RetriGen—a novel end-to-end framework that introduces a hybrid retriever jointly leveraging lexical (token-level) and semantic (embedding-level) retrieval, coupled with collaborative fine-tuning of a pre-trained language model (PLM) generator. This design enables cross-repository, semantics-aware assertion retrieval and precise generation. Evaluated on two established benchmarks, RetriGen achieves 57.66% assertion accuracy and 73.24% CodeBLEU—outperforming the strongest baseline by +50.66% and +14.14%, respectively. These results demonstrate substantial improvements in both the depth of semantic reasoning and the robustness of generated assertions.

Technology Category

• Natural Language Processing: Code Generation / Program Synthesis from Natural Language
• Data Mining & Knowledge Management: Conversational Systems for Recommendation & Retrieval
• Knowledge Representation and Reasoning: Diagnosis and Abductive Reasoning

Application Category

• Search and Retrieval-Augmented AI: Large language models for search
• Semantics and Knowledge: Methods to enhance, augment, integrate or synergize semantic models such as knowledge graphs and LLMs
• User Modeling, Personalization and Recommendation: Fairness-aware retrieval and ranking

Abstract

Unit testing validates the correctness of the units of the software system under test and serves as the cornerstone in improving software quality and reliability. To reduce manual efforts in writing unit tests, some techniques have been proposed to automatically generate test assertions, with recent integration-based approaches considered state-of-the-art. Despite being promising, such integration-based approaches face several limitations, including reliance on lexical matching for assertion retrieval and a limited training corpus for assertion generation. This paper proposes a novel retrieval-augmented deep assertion generation approach, namely RetriGen, based on a hybrid retriever and a pre-trained language model (PLM)-based generator. Given a focal-test, RetriGen first builds a hybrid assertion retriever to search for the most relevant Test-Assert Pair from external codebases. The retrieval process considers lexical similarity and semantical similarity via a token-based and an embedding-based retriever, respectively. RetriGen then treats assertion generation as a sequence-to-sequence task and designs a PLM-based assertion generator to predict a correct assertion. We conduct extensive experiments to evaluate RetriGen against six state-of-the-art approaches across two large-scale datasets and two metrics. The results demonstrate that RetriGen achieves 57.66% accuracy and 73.24% CodeBLEU, outperforming all baselines with average improvements of 50.66% and 14.14%, respectively.

Problem

• Automates unit test assertion generation
• Enhances retrieval with hybrid methods
• Improves accuracy using pre-trained models

Innovation

• Hybrid retriever for assertion retrieval
• Pre-trained language model generator
• Combines lexical and semantical similarity