AI-Assisted Unit Test Writing and Test-Driven Code Refactoring: A Case Study
| Authors | Ema Smolic et al. |
| Year | 2026 |
| Field | AI / ML |
| arXiv | 2604.03135 |
| Download | |
| Categories | cs.SE, cs.AI |
Abstract
Many software systems originate as prototypes or minimum viable products (MVPs), developed with an emphasis on delivery speed and responsiveness to changing requirements rather than long-term code maintainability. While effective for rapid delivery, this approach can result in codebases that are difficult to modify, presenting a significant opportunity cost in the era of AI-assisted or even AI-led programming. In this paper, we present a case study of using coding models for automated unit test generation and subsequent safe refactoring, with proposed code changes validated by passing tests. The study examines best practices for iteratively generating tests to capture existing system behavior, followed by model-assisted refactoring under developer supervision. We describe how this workflow constrained refactoring changes, the errors and limitations observed in both phases, the efficiency gains achieved, when manual intervention was necessary, and how we addressed the weak value misalignment we observed in models. Using this approach, we generated nearly 16,000 lines of reliable unit tests in hours rather than weeks, achieved up to 78% branch coverage in critical modules, and significantly reduced regression risk during large-scale refactoring. These results illustrate software engineering's shift toward an empirical science, emphasizing data collection and constraining mechanisms that support fast, safe iteration.
Engineering Breakdown
Plain English
This paper demonstrates a practical workflow for using AI coding models to automatically generate unit tests for legacy codebases and then safely refactor that code with the tests validating correctness. The authors conducted a case study showing how this two-stage approach—test generation followed by model-assisted refactoring under developer supervision—can improve code maintainability without breaking existing functionality. The core insight is that coding models can capture existing system behavior through test generation, which then serves as a safety net for subsequent refactoring. This addresses a real production problem: many teams have rapid-prototype codebases that are expensive to modify, and AI models can help bridge the gap between speed-to-market and long-term code quality.
Core Technical Contribution
The paper's primary contribution is a validated workflow pattern for applying coding models to legacy code improvement, not a new model architecture or training technique. The novelty lies in the iterative, two-phase methodology: first using models to generate tests that capture existing behavior (effectively reverse-engineering test coverage), then using those tests as validation gates for model-generated refactoring suggestions. This is distinct from prior work that either focuses solely on code generation or test generation in isolation; the authors show how chaining these two tasks with developer oversight creates a safe, practical refactoring pipeline. The case study validates that this pattern is feasible with current coding models and produces actionable improvements.
How It Works
The workflow operates in two sequential phases. First, developers feed the AI model existing source code along with specifications or context about what the code should do, and the model generates candidate unit tests that exercise the system under its current behavior. These generated tests are then executed against the codebase to ensure they pass, establishing a baseline test suite that captures the system's actual behavior rather than its intended behavior. In the second phase, the validated test suite becomes the evaluation metric: developers work with the AI model to propose refactoring changes (improving structure, removing duplication, modernizing patterns), and each proposed change is validated by re-running the test suite. Only refactoring suggestions that maintain test passage are accepted, giving developers confidence that behavioral contracts are preserved. The developer remains in the loop at each step, reviewing and filtering model outputs rather than accepting them blindly.
Production Impact
For teams managing legacy codebases, this approach dramatically reduces the friction and risk of refactoring—the biggest blocker to code modernization in many organizations. Instead of manually writing test coverage before refactoring (a high-effort task that gets deferred), teams can generate test baselines quickly with AI models, then incrementally improve code quality with quantifiable validation. In production pipelines, this could be integrated as an automated workflow: run test generation on a service, automatically execute and validate the tests, then periodically propose and validate refactoring suggestions that improve metrics like cyclomatic complexity or method length. Trade-offs include: significant latency per refactoring cycle (model inference + test execution), potential for generated tests to encode bugs or corner cases that shouldn't be preserved, and the need for human oversight to prevent subtle behavioral changes. Compute costs scale with codebase size, and you must have reliable test infrastructure for this to work.
Limitations and When Not to Use This
The paper assumes that generated tests accurately capture intended behavior, but generated tests can encode bugs, incomplete coverage, or overly specific assertions that make future refactoring harder. The approach is only as good as the coding model's understanding of the codebase—it struggles with implicit contracts, undocumented APIs, or domain-specific constraints not visible in the code. The case study is single-threaded: it's unclear whether this workflow scales to very large monoliths, polyglot codebases, or systems with external dependencies and integration test requirements that can't be captured in unit tests. The paper does not address how to handle tests that become obsolete or overly restrictive after refactoring, nor does it provide guidance on when to halt refactoring or how to measure improvement beyond test passage.
Research Context
This work sits at the intersection of two active research areas: AI-assisted code generation (models like Codex, GPT-4, and open-source alternatives) and automated testing/test synthesis. It builds on prior research in test generation (symbolic execution, property-based testing) but applies modern large language models as the generation mechanism, which is a shift in approach. The paper contributes to the emerging field of AI-driven software engineering practices, similar to recent work on AI-assisted code review, automated debugging, and model-guided refactoring. It opens up the research direction of studying long-horizon code improvement workflows where multiple AI tasks (testing, refactoring, documentation) are chained together with human validation, rather than treating each task in isolation.
:::tip Subscribe Get weekly breakdowns of papers like this in AI Letters - the newsletter for engineers building production AI systems. :::