使用Spring AI和LLM生成Java测试代码

背景

      AIDocumentLibraryChat 项目已扩展至生成测试代码(Java 代码已通过测试)。该项目可为公开的 Github 项目生成测试代码。只需提供要测试的类的网址,该类就会被加载、分析导入,项目中的依赖类也会被加载。这样,LLM 就有机会在为测试生成模拟时考虑导入的源类。可以提供 testUrl,为 LLM 提供一个示例,以便生成测试。我们已使用 Ollama 对 granite-codedeepseek-coder-v2codestral 模型进行了测试。

目的是测试 LLM 对开发人员创建测试的帮助有多大。

实施

配置


要选择 LLM 模型,需要更新 application-ollama.properties 文件

spring.ai.ollama.base-url=${OLLAMA-BASE-URL:http://localhost:11434}
spring.ai.ollama.embedding.enabled=false
spring.ai.embedding.transformer.enabled=true
document-token-limit=150
embedding-token-limit=500
spring.liquibase.change-log=classpath:/dbchangelog/db.changelog-master-ollama.xml

...

# generate code
#spring.ai.ollama.chat.model=granite-code:20b
#spring.ai.ollama.chat.options.num-ctx=8192

spring.ai.ollama.chat.options.num-thread=8
spring.ai.ollama.chat.options.keep_alive=1s

#spring.ai.ollama.chat.model=deepseek-coder-v2:16b
#spring.ai.ollama.chat.options.num-ctx=65536

spring.ai.ollama.chat.model=codestral:22b
spring.ai.ollama.chat.options.num-ctx=32768

spring.ai.olama.chat.model "选择要使用的 LLM 代码模型。

spring.ollama.chat.options.num-ctx "设置上下文窗口中的令牌数量。上下文窗口包含请求所需的令牌和响应所需的令牌。

如果 Ollama 没有选择正确的内核数量,可以使用 “spring.olama.chat.options.num-thread”。spring.olama.chat.options.keep_alive "设置了保留上下文窗口的秒数。

Controller


Controller是获取信号源和生成测试的接口:

@RestController
@RequestMapping("rest/code-generation")
public class CodeGenerationController {
   private final CodeGenerationService codeGenerationService;

  public CodeGenerationController(CodeGenerationService
     codeGenerationService) {
     this.codeGenerationService = codeGenerationService;
   }

  @GetMapping("/test")
   public String getGenerateTests(@RequestParam("url") String url,
     @RequestParam(name = "testUrl", required = false) String testUrl) {
     return this.codeGenerationService.generateTest(URLDecoder.decode(url,
       StandardCharsets.UTF_8),
     Optional.ofNullable(testUrl).map(myValue -> URLDecoder.decode(myValue,
       StandardCharsets.UTF_8)));
   }

  @GetMapping("/sources")
   public GithubSources getSources(@RequestParam("url") String url,
     @RequestParam(name="testUrl", required = false) String testUrl) {
     var sources = this.codeGenerationService.createTestSources(
       URLDecoder.decode(url, StandardCharsets.UTF_8), true);
     var test = Optional.ofNullable(testUrl).map(myTestUrl ->
       this.codeGenerationService.createTestSources(
         URLDecoder.decode(myTestUrl, StandardCharsets.UTF_8), false))
           .orElse(new GithubSource("none", "none", List.of(), List.of()));
     return new GithubSources(sources, test);
   }
}

代码生成控制器 “有一个 ”getSources(...) "方法。该方法获取要生成测试的类的 URL 和可选的 testUrl,以及可选的示例测试。它对请求参数进行解码,并调用 “createTestSources(...) ”方法。该方法会返回 “GithubSources”,其中包含要测试的类的源代码、其在项目中的依赖关系以及测试示例。

getGenerateTests(...) “方法获取测试类的 ”url “和可选的 ”testUrl “以进行 url 解码,并调用 ”CodeGenerationService “的 ”generateTests(...) "方法。

Services


CodeGenerationService "从 Github 收集类,并为被测类生成测试代码。

带有提示的服务看起来像这样:

@Service
public class CodeGenerationService {
   private static final Logger LOGGER = LoggerFactory
     .getLogger(CodeGenerationService.class);
   private final GithubClient githubClient;
   private final ChatClient chatClient;
   private final String ollamaPrompt = """
     You are an assistant to generate spring tests for the class under test.
     Analyse the classes provided and generate tests for all methods. Base 
     your tests on the example.
     Generate and implement the test methods. Generate and implement complete 
     tests methods.
     Generate the complete source of the test class.
                     
     Generate tests for this class:
     {classToTest}

    Use these classes as context for the tests:
     {contextClasses}

    {testExample}
   """;   
   private final String ollamaPrompt1 = """
     You are an assistant to generate a spring test class for the source
     class.
     1. Analyse the source class
     2. Analyse th