Test code generation at Ericsson using Program Analysis Augmented Fine Tuned LLMs

TL;DR

This paper presents a method for generating Java test code from natural language test steps at Ericsson, combining retrieval-augmented generation, static analysis, and fine-tuning of large language models to improve accuracy.

Contribution

The authors introduce a novel approach that integrates static program analysis with retrieval-augmented prompting and fine-tuning of LLMs for test code generation.

Findings

Fine-tuning improves code conformity by 8% in F1-score.

Retrieval-augmented prompts help mitigate function signature assumptions.

Fine-tuned 8x7b MoE model performs comparably to larger models.

Abstract

We describe test code generation using Large Language Models (LLMs) in Ericsson. Our input is a test step in natural language (English) and our output is code (Java) which accomplishes the test step. We describe how straight forward prompting does not suffice and results in LLM assuming functions and signatures which are not present in the code repository. We then show how we alleviate the problem by a combination of Retrieval Augmented Generation (RAG) along with prompt engineering that expanded the simple prompt with additional contextual information using static program analysis. We then describe further improvements that we obtained by fine-tuning the underlying LLM. The fine tuning is done based on a custom designed prompt template which has pre-dependent classes, their public methods as well two exemplar outputs obtained from RAG. Our results establish that our fine tuned models help improve the correspondence or conformity with the original developer written test code as measured by the traditional metrics of F1-score based on the methods used in the generated code. Fine tuning of a 8x7b Mixture of Experts (MoE) model leads to an average improvement of 8\% over the base model and is comparable to the scores on a much larger 8x22b MoE model.