Building a Test-Generation Agent with Local SLMs
Automated test generation has long been a holy grail for software teams. With the rise of Small Language Models (SLMs) — compact models in the 1B to 8B parameter range — it is now practical to run a capable test-generation agent entirely on a developer laptop or CI runner, with no external API calls, no data leaving the machine, and zero per-token cost. This tutorial walks through the architecture, implementation, and best practices for building such an agent from scratch.
What Is a Test-Generation Agent?
A test-generation agent is an autonomous loop that reads source code, reasons about its behavior, and produces executable tests — typically unit tests — that exercise meaningful paths through the code. Unlike a one-shot "generate tests for this function" prompt, an agent iterates: it inspects the code, drafts tests, runs them, observes failures, and refines until the tests pass and provide real coverage.
Running this agent on a local SLM means the model is loaded into memory on the same machine doing the work. Popular choices include Qwen2.5-Coder (7B), Phi-3.5-mini, Llama-3.1-8B, and DeepSeek-Coder. These models are small enough to run on consumer GPUs or even CPU-only machines via quantized formats like GGUF.
Why Local SLMs Matter
- Privacy and compliance: Source code never leaves the machine, which matters for proprietary codebases and regulated industries.
- Zero marginal cost: After the one-time model download, you can generate unlimited tests without API bills.
- Low latency: No network round-trips means the agent loop can iterate quickly.
- Offline operation: Works on air-gapped machines and during travel.
- Reproducibility: A pinned model version produces deterministic-ish outputs, which is critical for CI integration.
- CI friendliness: Self-hosted runners can execute the agent as a pre-merge step without exposing code to third parties.
The trade-off is raw capability: an SLM will not match GPT-4-class models on complex reasoning. The key is to compensate with a tight agent loop, good prompting, and verification through actual test execution.
Architecture of the Agent
The agent is built around a small number of components:
- Model server: Loads the SLM and exposes a completion API. We will use
llama.cppvia its OpenAI-compatible server. - Code indexer: Extracts functions, classes, and signatures from the target source files.
- Test drafter: Prompts the model to produce candidate tests for a given symbol.
- Test runner: Writes the candidate tests to disk and executes the project's test command, capturing stdout, stderr, and exit code.
- Refinement loop: Feeds failures back to the model for a fixed number of iterations.
- Orchestrator: Ties everything together and persists accepted tests.
Setting Up the Local Model Server
Install llama.cpp and download a quantized coder model. The GGUF format is the most portable option.
# Clone and build llama.cpp
git clone https://github.com/ggerganov/llama.cpp
cd llama.cpp
make LLAMA_CUDA=1 # omit LLAMA_CUDA for CPU-only
# Download a quantized coder model
huggingface-cli download Qwen/Qwen2.5-Coder-7B-Instruct-GGUF \
qwen2.5-coder-7b-instruct-q5_k_m.gguf \
--local-dir ./models
# Start the OpenAI-compatible server
./llama-server \
-m ./models/qwen2.5-coder-7b-instruct-q5_k_m.gguf \
--port 8080 \
-c 8192 \
-ngl 33
The server now exposes http://localhost:8080/v1/chat/completions, which is API-compatible with the OpenAI Python client. This means you can swap to a cloud model later by changing a single base URL.
Implementing the Agent
We will build the agent in Python. Install the dependencies:
pip install openai tiktoken tree-sitter tree-sitter-python
1. The Model Client Wrapper
First, a thin wrapper around the OpenAI client pointed at our local server:
from openai import OpenAI
from typing import List, Dict
client = OpenAI(base_url="http://localhost:8080/v1", api_key="not-needed")
MODEL = "qwen2.5-coder-7b-instruct"
def chat(messages: List[Dict], temperature: float = 0.2, max_tokens: int = 2048) -> str:
response = client.chat.completions.create(
model=MODEL,
messages=messages,
temperature=temperature,
max_tokens=max_tokens,
stop=["</TEST_FILE>"],
)
return response.choices[0].message.content
Low temperature is important for code generation: we want deterministic, syntactically valid output rather than creative variation.
2. Extracting Symbols from Source
Use tree-sitter to parse the target file and pull out top-level functions and classes. This gives the agent focused, well-bounded units to test.
from tree_sitter import Parser
from tree_sitter_python import language as python_language
parser = Parser(python_language)
def extract_symbols(source: str) -> List[dict]:
tree = parser.parse(bytes(source, "utf8"))
symbols = []
def walk(node, indent=0):
if node.type in ("function_definition", "class_definition"):
name_node = node.child_by_field_name("name")
if name_node:
symbols.append({
"name": name_node.text.decode("utf8"),
"type": node.type,
"start": node.start_point[0],
"end": node.end_point[0],
"body": source.splitlines()[node.start_point[0]:node.end_point[0]+1],
})
for child in node.children:
walk(child, indent+1)
walk(tree.root_node)
return symbols
3. The Test Drafter
The drafter builds a prompt that includes the source code, the project's testing conventions, and a strict output format. Constraining the output format is the single most effective technique for getting usable results from an SLM.
DRAFT_SYSTEM = """You are a test-generation agent for Python projects.
You write pytest unit tests. You output ONLY the test file content
wrapped in <TEST_FILE> tags. No prose, no explanations."""
DRAFT_USER = """Generate pytest unit tests for the following code.
Requirements:
- Use pytest style with `def test_*` functions.
- Cover normal cases, edge cases, and error cases.
- Import the target from the module path shown.
- Do NOT mock the target function itself.
- Keep tests deterministic (no random data without a seed).
Module path: {module}
Source file: {filename}
python
{source}
Output the test file inside <TEST_FILE>...</TEST_FILE> tags."""
def draft_tests(module: str, filename: str, source: str) -> str:
messages = [
{"role": "system", "content": DRAFT_SYSTEM},
{"role": "user", "content": DRAFT_USER.format(
module=module, filename=filename, source=source
)},
]
raw = chat(messages, temperature=0.2, max_tokens=3000)
# Extract content between tags
if "<TEST_FILE>" in raw:
return raw.split("<TEST_FILE>")[1].split("</TEST_FILE>")[0].strip()
return raw.strip()
4. The Test Runner
The runner writes the candidate test to a file and executes pytest in a subprocess, capturing the result. This is the verification step that turns a language model into an agent.
import subprocess
import tempfile
import os
import pathlib
def run_tests(test_code: str, project_root: str, test_path: str) -> dict:
test_full = os.path.join(project_root, test_path)
pathlib.Path(test_full).parent.mkdir(parents=True, exist_ok=True)
with open(test_full, "w") as f:
f.write(test_code)
result = subprocess.run(
["python", "-m", "pytest", test_full, "-x", "--tb=short", "-q"],
capture_output=True,
text=True,
cwd=project_root,
timeout=120,
)
return {
"exit_code": result.returncode,
"stdout": result.stdout,
"stderr": result.stderr,
"passed": result.returncode == 0,
}
5. The Refinement Loop
When tests fail, the agent feeds the failure output back to the model and asks for a corrected version. This is where the agent behavior lives.
REFINE_USER = """The test file you generated failed. Fix it.
Original source:
python
{source}
Test file:
python
{test_code}
Pytest output:
{output}
Output the corrected test file inside <TEST_FILE>...</TEST_FILE> tags."""
def refine_tests(source: str, test_code: str, output: str) -> str:
messages = [
{"role": "system", "content": DRAFT_SYSTEM},
{"role": "user", "content": REFINE_USER.format(
source=source, test_code=test_code, output=output[-3000:]
)},
]
raw = chat(messages, temperature=0.1, max_tokens=3000)
if "<TEST_FILE>" in raw:
return raw.split("<TEST_FILE>")[1].split("</TEST_FILE>")[0].strip()
return raw.strip()
6. The Orchestrator
The orchestrator ties the pieces together. For each source file, it extracts symbols, drafts tests, runs them, and refines up to a maximum number of iterations.
def generate_tests_for_file(source_path: str, project_root: str,
module: str, max_iterations: int = 3) -> dict:
with open(source_path) as f:
source = f.read()
test_rel = source_path.replace("src/", "tests/test_").replace(".py", "_test.py")
print(f"Drafting tests for {module}...")
test_code = draft_tests(module, source_path, source)
for i in range(max_iterations):
result = run_tests(test_code, project_root, test_rel)
if result["passed"]:
print(f"Tests passed on iteration {i+1}.")
return {"success": True, "iterations": i + 1,
"test_path": test_rel, "test_code": test_code}
print(f"Iteration {i+1} failed. Refining...")
combined_output = result["stdout"] + "\n" + result["stderr"]
test_code = refine_tests(source, test_code, combined_output)
# Final run after last refinement
result = run_tests(test_code, project_root, test_rel)
return {
"success": result["passed"],
"iterations": max_iterations,
"test_path": test_rel,
"test_code": test_code,
"final_output": result["stdout"],
}
Running it against a project is then a simple loop:
import glob
project_root = "/home/dev/myproject"
source_files = glob.glob(f"{project_root}/src/**/*.py", recursive=True)
results = []
for path in source_files:
module = path.replace(project_root + "/", "").replace("/", ".").replace(".py", "")
results.append(generate_tests_for_file(path, project_root, module))
passed = sum(1 for r in results if r["success"])
print(f"Generated passing tests for {passed}/{len(results)} files.")
Best Practices
Constrain the Output Relentlessly
SLMs drift easily. Use explicit output tags, provide a concrete example in the system prompt, and parse defensively. The <TEST_FILE> tag pattern lets you recover the code block even when the model adds stray commentary.
Keep Context Windows Tight
A 7B model with an 8K context window cannot ingest an entire repository. Feed it one symbol or one file at a time. If a file is large, extract only the target function plus its immediate dependencies rather than the whole module.
Provide Project Conventions Up Front
Include the existing test style, import patterns, fixture conventions, and the test runner command in the system prompt. A few-shot example drawn from the project's own passing tests dramatically improves output quality.
FEW_SHOT = """
Example of a good test from this project:
<TEST_FILE>
from myproject.math_utils import clamp
def test_clamp_within_bounds():
assert clamp(5, 0, 10) == 5
def test_clamp_below_min():
assert clamp(-3, 0, 10) == 0
def test_clamp_above_max():
assert clamp(99, 0, 10) == 10
</TEST_FILE>
"""
Always Verify by Execution
Never trust generated tests without running them. A test that imports a non-existent symbol or calls a function with the wrong signature is worse than no test because it creates false confidence. The run-then-refine loop is non-negotiable.
Guard Against Trivial Tests
SLMs sometimes produce tests that pass but assert nothing meaningful — for example, calling a function and checking it does not raise. Add a post-processing step that rejects tests with no assert statements or with only assert True.
def is_meaningful(test_code: str) -> bool:
lines = [l.strip() for l in test_code.splitlines()]
assert_lines = [l for l in lines if l.startswith("assert ") and "assert True" not in l]
return len(assert_lines) >= 2
Pin the Model and Seed
In CI, pin the exact model file and set temperature=0 (or a fixed seed if your server supports it). This makes test generation reproducible across runs, which is essential when the agent runs as a pre-merge check.
Respect Existing Coverage
Before generating tests, run a coverage tool and skip symbols that are already well covered. This prevents the agent from producing redundant tests and focuses its limited capability on genuine gaps.
Use Smaller Models for Triage
A 1.5B model is fast enough to classify symbols as "trivial," "needs tests," or "too complex for auto-generation." Use it as a pre-filter so the larger 7B model only spends tokens on symbols worth testing.
Integrating with CI
To run the agent in CI, wrap the orchestrator in a script that exits non-zero if the success rate falls below a threshold. This prevents regressions in test generation from silently shipping broken tests.
# ci_scripts/gen_tests.py
import sys
results = run_orchestrator()
success_rate = sum(1 for r in results if r["success"]) / len(results)
print(f"Success rate: {success_rate:.0%}")
if success_rate < 0.8:
print("Test generation success rate below threshold.")
sys.exit(1)
Run this on a self-hosted runner with a GPU, or on a CPU-only runner with a small Q4 quantized model. A Qwen2.5-Coder-7B Q4 model generates roughly 20-40 tokens per second on a modern CPU, which is fast enough for an overnight or pre-merge job.
Conclusion
Building a test-generation agent with local SLMs is now a practical, repeatable exercise. The combination of a constrained output format, a tight draft-run-refine loop, and execution-based verification compensates for the smaller model's reduced reasoning capacity. By keeping everything local, you gain privacy, cost predictability, and CI reproducibility — qualities that matter more than peak benchmark scores for day-to-day developer tooling. Start with a single well-tested module, tune the prompts to your project's conventions, and expand coverage incrementally. The result is an agent that earns its place in your development workflow by producing tests that actually pass and actually test something.