CAT-LM: Training Language Models on Aligned Code And Tests

TL;DR

CAT-LM is a large language model trained on aligned code and test data, leveraging a novel pretraining approach and extended context window to generate high-quality, developer-like tests that achieve coverage efficiently.

Contribution

The paper introduces CAT-LM, a GPT-style model trained on aligned code and test data with a new pretraining signal and longer context, improving test generation quality and relevance.

Findings

CAT-LM produces tests with coverage similar to human-written tests.

It outperforms larger models like CodeGen 16B and StarCoder in test validity.

Sampling with filtering enhances test quality and applicability.

Abstract

Testing is an integral part of the software development process. Yet, writing tests is time-consuming and therefore often neglected. Classical test generation tools such as EvoSuite generate behavioral test suites by optimizing for coverage, but tend to produce tests that are hard to understand. Language models trained on code can generate code that is highly similar to that written by humans, but current models are trained to generate each file separately, as is standard practice in natural language processing, and thus fail to consider the code-under-test context when producing a test file. In this work, we propose the Aligned Code And Tests Language Model (CAT-LM), a GPT-style language model with 2.7 Billion parameters, trained on a corpus of Python and Java projects. We utilize a novel pretraining signal that explicitly considers the mapping between code and test files when available. We also drastically increase the maximum sequence length of inputs to 8,192 tokens, 4x more than typical code generation models, to ensure that the code context is available to the model when generating test code. We analyze its usefulness for realistic applications, showing that sampling with filtering (e.g., by compilability, coverage) allows it to efficiently produce tests that achieve coverage similar to ones written by developers while resembling their writing style. By utilizing the code context, CAT-LM generates more valid tests than even much larger language models trained with more data (CodeGen 16B and StarCoder) and substantially outperforms a recent test-specific model (TeCo) at test completion. Overall, our work highlights the importance of incorporating software-specific insights when training language models for code and paves the way to more powerful automated test generation.