INTRODUCTION AND MOTIVATION
Creating comprehensive unit tests is one of the most time-consuming yet critical aspects of software development. Developers must consider not only the happy path scenarios but also edge cases, boundary conditions, error handling, and various input combinations. A unit test generator powered by Large Language Models offers an innovative solution to this challenge by automatically analyzing source code and producing thorough test suites that achieve high code coverage while addressing potential edge cases.
The fundamental value proposition of an LLM-based test generator lies in its ability to understand code semantics beyond simple pattern matching. Traditional static analysis tools can identify code paths and branches, but they struggle to generate meaningful test cases that reflect real-world usage scenarios. Large Language Models, trained on vast repositories of code and associated tests, can infer the intent behind code structures and generate tests that not only exercise the code but also validate expected behaviors in contextually appropriate ways.
Modern software projects often contain large files with thousands of lines of code, presenting a significant challenge for LLM-based analysis. Language models have finite context windows, typically ranging from four thousand to one hundred thousand tokens. A large implementation file can easily exceed these limits, making it impossible to process the entire file in a single LLM call. This necessitates a Retrieval-Augmented Generation approach where the code is intelligently chunked, embedded into vector representations, stored in a vector database, and selectively retrieved based on semantic relevance to the current test generation task.
This article presents a comprehensive architecture for building such a system. The generator must be flexible enough to handle multiple programming languages, support both local and remote LLM deployments, leverage various GPU architectures for optimal performance, and employ sophisticated RAG techniques for processing large codebases. The system analyzes implementation files, extracts relevant code structures through semantic retrieval, generates appropriate test cases, and ensures comprehensive coverage including edge case handling.
ARCHITECTURAL OVERVIEW
The unit test generator consists of several interconnected components that work together to transform source code into comprehensive test suites. The architecture follows clean architecture principles with clear separation of concerns and well-defined interfaces between layers.
The first major component is the Code Analysis Layer, which is responsible for parsing source files and extracting structural information. This layer must handle multiple programming languages, each with its own syntax and semantics. The analyzer identifies functions, methods, classes, parameters, return types, and dependencies. It also performs control flow analysis to understand branching logic and identify code paths that require testing. When dealing with large files, this layer works in conjunction with the RAG system to process code in manageable chunks.
The second component is the Code Chunking and Embedding Layer, which becomes critical when processing large files that exceed LLM context windows. Unlike naive text splitting approaches that arbitrarily divide code at character or line boundaries, this layer employs code-aware chunking strategies that respect syntactic and semantic boundaries. The chunker identifies natural division points such as function boundaries, class definitions, and logical code blocks. Each chunk is then converted into a dense vector embedding that captures its semantic meaning, allowing for similarity-based retrieval during test generation.
The third component is the Vector Database Integration Layer, which stores code chunk embeddings and enables efficient semantic search. This layer provides a unified interface to various vector database backends including Chroma, Pinecone, Weaviate, and FAISS. The database stores not only the embeddings but also metadata about each chunk including its location in the original file, dependencies, and structural information. During test generation, the system queries this database to retrieve the most relevant code chunks for the function or class being tested.
The fourth component is the LLM Integration Layer, which provides a unified interface for interacting with both local and remote language models. This abstraction allows the system to seamlessly switch between different LLM providers without affecting other components. The integration layer handles prompt construction, response parsing, token management, and error handling. It also manages GPU acceleration when using local models, detecting available hardware and configuring the appropriate backend for Intel, AMD ROCm, Apple MPS, or Nvidia CUDA architectures.
The fifth component is the Test Generation Engine, which orchestrates the entire test creation process. This engine combines insights from the code analyzer with the generative capabilities of the LLM to produce test cases. When working with large files, it uses the RAG system to retrieve relevant context before generating each test. It maintains templates for different testing frameworks, manages test case organization, and ensures that generated tests follow best practices for the target language and framework.
The sixth component is the Coverage Analysis Module, which evaluates the generated tests to ensure they adequately exercise the code under test. This module performs static analysis on both the implementation and the test code to identify uncovered branches, paths, and edge cases. When gaps are detected, it triggers additional test generation iterations to fill those gaps, using the RAG system to retrieve additional context about uncovered code sections.
The seventh component is the Edge Case Detector, which uses both static analysis and LLM reasoning to identify potential edge cases that require explicit testing. This includes boundary values, null or empty inputs, type mismatches, concurrent access scenarios, and exceptional conditions. The detector maintains a knowledge base of common edge cases for different data types and programming patterns.
CODE CHUNKING STRATEGIES FOR RETRIEVAL-AUGMENTED GENERATION
The effectiveness of the RAG approach depends critically on how code is divided into chunks. Naive chunking strategies that split code at arbitrary character counts or line numbers inevitably break syntactic structures, severing the relationship between related code elements and producing fragments that lack semantic coherence. A code-aware chunking strategy must respect the hierarchical structure of source code and preserve the semantic relationships that enable meaningful test generation.
The fundamental principle of code-aware chunking is to identify natural boundaries in the code structure. In object-oriented languages, classes represent natural chunk boundaries because they encapsulate related functionality and data. Within a class, individual methods form logical units that can be chunked independently while maintaining references to their containing class context. In functional programming languages, top-level functions and their associated helper functions form natural groupings.
The chunking strategy must also consider dependencies and relationships between code elements. A method that calls several helper methods should ideally be chunked together with those helpers, or at minimum, the chunk metadata should record these dependencies so the retrieval system can fetch related chunks when needed. Similarly, a class that inherits from a base class or implements interfaces should maintain references to those parent structures.
Here is an example of how the code chunker analyzes and divides Python code while respecting structural boundaries:
import ast
from typing import List, Dict, Any, Optional, Set
from dataclasses import dataclass
@dataclass
class CodeChunk:
content: str
chunk_type: str # 'function', 'class', 'method', 'module'
name: str
start_line: int
end_line: int
dependencies: Set[str]
parent_context: Optional[str]
metadata: Dict[str, Any]
def get_full_context(self) -> str:
"""Returns chunk content with parent context if available."""
if self.parent_context:
return f"{self.parent_context}\n\n{self.content}"
return self.content
class CodeAwareChunker:
def __init__(self, source_code: str, language: str, max_chunk_size: int = 1000):
self.source_code = source_code
self.language = language
self.max_chunk_size = max_chunk_size
self.source_lines = source_code.split('\n')
def chunk_code(self) -> List[CodeChunk]:
"""Main entry point for chunking code based on language."""
if self.language == 'python':
return self._chunk_python()
elif self.language == 'javascript':
return self._chunk_javascript()
elif self.language == 'java':
return self._chunk_java()
else:
return self._chunk_generic()
def _chunk_python(self) -> List[CodeChunk]:
"""Chunks Python code respecting AST structure."""
try:
tree = ast.parse(self.source_code)
except SyntaxError as e:
# Fall back to line-based chunking if parsing fails
return self._chunk_generic()
chunks = []
module_imports = self._extract_imports(tree)
for node in ast.iter_child_nodes(tree):
if isinstance(node, ast.ClassDef):
class_chunks = self._chunk_class(node, module_imports)
chunks.extend(class_chunks)
elif isinstance(node, ast.FunctionDef):
func_chunk = self._chunk_function(node, module_imports, None)
chunks.append(func_chunk)
elif isinstance(node, (ast.Import, ast.ImportFrom)):
continue # Already captured in module_imports
else:
# Module-level code
chunk = self._create_module_level_chunk(node, module_imports)
if chunk:
chunks.append(chunk)
return chunks
def _chunk_class(self, node: ast.ClassDef, imports: str) -> List[CodeChunk]:
"""Chunks a class, potentially splitting large classes into multiple chunks."""
class_header = self._extract_class_header(node)
class_start = node.lineno - 1
class_end = node.end_lineno
chunks = []
methods = [n for n in node.body if isinstance(n, ast.FunctionDef)]
# Calculate class size
class_size = class_end - class_start
if class_size <= self.max_chunk_size:
# Small class - chunk as single unit
content = '\n'.join(self.source_lines[class_start:class_end])
full_content = f"{imports}\n\n{content}" if imports else content
dependencies = self._extract_dependencies(node)
chunk = CodeChunk(
content=full_content,
chunk_type='class',
name=node.name,
start_line=class_start,
end_line=class_end,
dependencies=dependencies,
parent_context=None,
metadata={
'bases': [self._get_name(b) for b in node.bases],
'decorators': [self._get_name(d) for d in node.decorator_list],
'method_count': len(methods)
}
)
chunks.append(chunk)
else:
# Large class - chunk methods individually with class context
for method in methods:
method_chunk = self._chunk_function(
method,
imports,
class_header
)
method_chunk.metadata['class_name'] = node.name
chunks.append(method_chunk)
return chunks
def _chunk_function(self, node: ast.FunctionDef, imports: str,
class_context: Optional[str]) -> CodeChunk:
"""Creates a chunk for a function or method."""
start_line = node.lineno - 1
end_line = node.end_lineno
content = '\n'.join(self.source_lines[start_line:end_line])
dependencies = self._extract_dependencies(node)
# Build full content with context
parts = []
if imports:
parts.append(imports)
if class_context:
parts.append(class_context)
parts.append(content)
full_content = '\n\n'.join(parts)
return CodeChunk(
content=full_content,
chunk_type='method' if class_context else 'function',
name=node.name,
start_line=start_line,
end_line=end_line,
dependencies=dependencies,
parent_context=class_context,
metadata={
'parameters': [arg.arg for arg in node.args.args],
'decorators': [self._get_name(d) for d in node.decorator_list],
'has_docstring': ast.get_docstring(node) is not None
}
)
def _extract_imports(self, tree: ast.AST) -> str:
"""Extracts all import statements from the module."""
import_lines = []
for node in ast.iter_child_nodes(tree):
if isinstance(node, (ast.Import, ast.ImportFrom)):
start = node.lineno - 1
end = node.end_lineno
import_lines.extend(self.source_lines[start:end])
return '\n'.join(import_lines)
def _extract_class_header(self, node: ast.ClassDef) -> str:
"""Extracts class definition line and docstring."""
start = node.lineno - 1
# Find the first method or end of class
first_method_line = None
for item in node.body:
if isinstance(item, ast.FunctionDef):
first_method_line = item.lineno - 1
break
if first_method_line:
end = first_method_line
else:
end = node.end_lineno
# Include class definition and docstring
header_lines = []
for i, line in enumerate(self.source_lines[start:end]):
header_lines.append(line)
# Stop after docstring if present
if i > 0 and '"""' in line or "'''" in line:
if line.count('"""') == 2 or line.count("'''") == 2:
break
elif line.strip().endswith('"""') or line.strip().endswith("'''"):
break
return '\n'.join(header_lines)
def _extract_dependencies(self, node: ast.AST) -> Set[str]:
"""Extracts function and class names referenced in the node."""
dependencies = set()
for child in ast.walk(node):
if isinstance(child, ast.Call):
func_name = self._get_name(child.func)
if func_name:
dependencies.add(func_name)
elif isinstance(child, ast.Name):
if isinstance(child.ctx, ast.Load):
dependencies.add(child.id)
return dependencies
def _get_name(self, node: ast.AST) -> Optional[str]:
"""Extracts name from various AST node types."""
if isinstance(node, ast.Name):
return node.id
elif isinstance(node, ast.Attribute):
value = self._get_name(node.value)
return f"{value}.{node.attr}" if value else node.attr
elif isinstance(node, ast.Call):
return self._get_name(node.func)
return None
def _create_module_level_chunk(self, node: ast.AST, imports: str) -> Optional[CodeChunk]:
"""Creates chunk for module-level code that isn't a function or class."""
if not hasattr(node, 'lineno'):
return None
start = node.lineno - 1
end = node.end_lineno
content = '\n'.join(self.source_lines[start:end])
return CodeChunk(
content=f"{imports}\n\n{content}" if imports else content,
chunk_type='module',
name='module_level',
start_line=start,
end_line=end,
dependencies=set(),
parent_context=None,
metadata={'node_type': type(node).__name__}
)
def _chunk_generic(self) -> List[CodeChunk]:
"""Fallback chunking strategy for unparseable or unsupported languages."""
chunks = []
current_start = 0
while current_start < len(self.source_lines):
end = min(current_start + self.max_chunk_size, len(self.source_lines))
content = '\n'.join(self.source_lines[current_start:end])
chunk = CodeChunk(
content=content,
chunk_type='generic',
name=f'chunk_{current_start}_{end}',
start_line=current_start,
end_line=end,
dependencies=set(),
parent_context=None,
metadata={}
)
chunks.append(chunk)
current_start = end
return chunks
The chunking strategy shown above demonstrates several important principles. First, it respects the Abstract Syntax Tree structure of the code, ensuring that functions and classes are never split mid-definition. Second, it preserves context by including import statements and class headers with each chunk, allowing the LLM to understand the chunk even when retrieved in isolation. Third, it extracts dependency information that can be used during retrieval to fetch related chunks. Fourth, it includes a fallback mechanism for code that cannot be parsed or for unsupported languages.
The maximum chunk size parameter requires careful tuning based on the target LLM's context window and the desired granularity of retrieval. Smaller chunks enable more precise retrieval but may lack sufficient context for the LLM to generate meaningful tests. Larger chunks provide more context but reduce retrieval precision and may still exceed context limits when multiple chunks are retrieved together. A typical value ranges from five hundred to fifteen hundred lines of code per chunk.
EMBEDDING GENERATION AND VECTOR DATABASE INTEGRATION
Once code has been chunked into semantically coherent units, each chunk must be converted into a dense vector embedding that captures its semantic meaning. These embeddings enable similarity-based retrieval where the system can find code chunks that are semantically related to a query, even if they do not share exact keywords or identifiers.
The embedding generation process uses specialized models trained on code rather than general-purpose text embeddings. Code-specific embedding models understand programming language syntax, common coding patterns, and the semantic relationships between code elements. Popular choices include CodeBERT, GraphCodeBERT, and UniXcoder, all of which have been pre-trained on large corpora of source code from multiple programming languages.
The embedding model processes each code chunk and produces a fixed-length vector, typically ranging from three hundred eighty-four to seven hundred sixty-eight dimensions. These vectors are designed so that semantically similar code chunks produce vectors that are close together in the embedding space, as measured by cosine similarity or Euclidean distance. For example, two functions that perform similar operations but use different variable names would produce similar embeddings.
Here is an example of the embedding generation and vector database integration:
from typing import List, Dict, Any, Optional, Tuple
import numpy as np
from sentence_transformers import SentenceTransformer
import chromadb
from chromadb.config import Settings
import torch
class CodeEmbeddingGenerator:
def __init__(self, model_name: str = 'microsoft/codebert-base',
device: Optional[str] = None):
"""
Initializes the embedding generator with a code-specific model.
Args:
model_name: Name of the embedding model to use
device: Device to run the model on (cuda, mps, cpu, or None for auto-detect)
"""
self.device = self._detect_device(device)
self.model = SentenceTransformer(model_name, device=self.device)
self.embedding_dimension = self.model.get_sentence_embedding_dimension()
def _detect_device(self, device: Optional[str]) -> str:
"""Detects the best available device for running the model."""
if device:
return device
# Check for CUDA (Nvidia)
if torch.cuda.is_available():
return 'cuda'
# Check for MPS (Apple Silicon)
if hasattr(torch.backends, 'mps') and torch.backends.mps.is_available():
return 'mps'
# Check for ROCm (AMD) - PyTorch treats it as CUDA
if torch.cuda.is_available() and 'rocm' in torch.version.hip:
return 'cuda' # ROCm uses CUDA API
# Default to CPU
return 'cpu'
def generate_embeddings(self, chunks: List[CodeChunk]) -> List[np.ndarray]:
"""
Generates embeddings for a list of code chunks.
Args:
chunks: List of CodeChunk objects
Returns:
List of embedding vectors as numpy arrays
"""
# Extract text content from chunks
texts = [chunk.get_full_context() for chunk in chunks]
# Generate embeddings in batches for efficiency
embeddings = self.model.encode(
texts,
batch_size=32,
show_progress_bar=True,
convert_to_numpy=True
)
return embeddings
def generate_embedding(self, text: str) -> np.ndarray:
"""Generates embedding for a single text."""
return self.model.encode(text, convert_to_numpy=True)
class VectorDatabaseManager:
def __init__(self, db_path: str = './chroma_db', collection_name: str = 'code_chunks'):
"""
Initializes the vector database manager.
Args:
db_path: Path to store the Chroma database
collection_name: Name of the collection to store chunks
"""
self.client = chromadb.Client(Settings(
chroma_db_impl="duckdb+parquet",
persist_directory=db_path
))
# Create or get collection
self.collection = self.client.get_or_create_collection(
name=collection_name,
metadata={"hnsw:space": "cosine"}
)
def store_chunks(self, chunks: List[CodeChunk], embeddings: List[np.ndarray]):
"""
Stores code chunks and their embeddings in the vector database.
Args:
chunks: List of CodeChunk objects
embeddings: List of embedding vectors
"""
# Prepare data for insertion
ids = [f"{chunk.name}_{chunk.start_line}_{chunk.end_line}" for chunk in chunks]
documents = [chunk.content for chunk in chunks]
metadatas = [self._prepare_metadata(chunk) for chunk in chunks]
# Convert numpy arrays to lists for Chroma
embeddings_list = [emb.tolist() for emb in embeddings]
# Store in batches
batch_size = 100
for i in range(0, len(chunks), batch_size):
batch_end = min(i + batch_size, len(chunks))
self.collection.add(
ids=ids[i:batch_end],
embeddings=embeddings_list[i:batch_end],
documents=documents[i:batch_end],
metadatas=metadatas[i:batch_end]
)
def _prepare_metadata(self, chunk: CodeChunk) -> Dict[str, Any]:
"""Prepares metadata for storage in vector database."""
metadata = {
'chunk_type': chunk.chunk_type,
'name': chunk.name,
'start_line': chunk.start_line,
'end_line': chunk.end_line,
'dependencies': ','.join(chunk.dependencies),
'has_parent_context': chunk.parent_context is not None
}
# Add custom metadata
for key, value in chunk.metadata.items():
if isinstance(value, (str, int, float, bool)):
metadata[key] = value
elif isinstance(value, list):
metadata[key] = ','.join(str(v) for v in value)
return metadata
def retrieve_similar_chunks(self, query_embedding: np.ndarray,
n_results: int = 5,
filter_metadata: Optional[Dict[str, Any]] = None) -> List[Dict[str, Any]]:
"""
Retrieves the most similar chunks to a query embedding.
Args:
query_embedding: Query vector
n_results: Number of results to return
filter_metadata: Optional metadata filters
Returns:
List of dictionaries containing chunk information
"""
results = self.collection.query(
query_embeddings=[query_embedding.tolist()],
n_results=n_results,
where=filter_metadata
)
# Format results
formatted_results = []
for i in range(len(results['ids'][0])):
formatted_results.append({
'id': results['ids'][0][i],
'content': results['documents'][0][i],
'metadata': results['metadatas'][0][i],
'distance': results['distances'][0][i] if 'distances' in results else None
})
return formatted_results
def retrieve_by_name(self, name: str) -> List[Dict[str, Any]]:
"""Retrieves chunks by exact name match."""
results = self.collection.get(
where={"name": name}
)
formatted_results = []
for i in range(len(results['ids'])):
formatted_results.append({
'id': results['ids'][i],
'content': results['documents'][i],
'metadata': results['metadatas'][i]
})
return formatted_results
def retrieve_dependencies(self, chunk_name: str, max_depth: int = 2) -> List[Dict[str, Any]]:
"""
Retrieves a chunk and its dependencies up to a specified depth.
Args:
chunk_name: Name of the chunk to start from
max_depth: Maximum dependency depth to traverse
Returns:
List of chunks including the original and dependencies
"""
visited = set()
results = []
def _retrieve_recursive(name: str, depth: int):
if depth > max_depth or name in visited:
return
visited.add(name)
chunks = self.retrieve_by_name(name)
for chunk in chunks:
results.append(chunk)
# Get dependencies from metadata
deps_str = chunk['metadata'].get('dependencies', '')
if deps_str:
dependencies = deps_str.split(',')
for dep in dependencies:
dep = dep.strip()
if dep and dep not in visited:
_retrieve_recursive(dep, depth + 1)
_retrieve_recursive(chunk_name, 0)
return results
def persist(self):
"""Persists the database to disk."""
self.client.persist()
The embedding generator shown above demonstrates how to leverage code-specific models while supporting multiple GPU architectures. The device detection logic checks for CUDA support which covers both Nvidia and AMD ROCm GPUs since PyTorch treats ROCm as a CUDA backend. It also checks for Apple's Metal Performance Shaders framework which accelerates operations on Apple Silicon. This ensures that the embedding generation process can take advantage of available hardware acceleration regardless of the underlying GPU architecture.
The vector database manager provides a clean interface to Chroma, a popular open-source vector database. The implementation includes methods for storing chunks with their embeddings, retrieving similar chunks based on semantic similarity, and traversing dependency graphs. The dependency retrieval method is particularly important for test generation because it allows the system to fetch not just the function being tested but also any helper functions or classes it depends on.
LLM INTEGRATION WITH MULTI-GPU SUPPORT
The LLM integration layer provides a unified interface for interacting with language models regardless of whether they are hosted remotely or running locally. This abstraction is critical because it allows the test generation system to work with various LLM providers without requiring changes to the core logic. The integration layer must handle prompt construction, response parsing, error handling, retry logic, and GPU acceleration for local models.
For remote LLMs, the integration typically involves making HTTP requests to API endpoints provided by services like OpenAI, Anthropic, or Cohere. These services handle all infrastructure concerns including model hosting, load balancing, and GPU allocation. The integration layer must manage API keys, rate limiting, token counting, and cost tracking.
For local LLMs, the integration is more complex because the system must load the model into memory, manage GPU resources, and handle inference directly. Local models offer several advantages including data privacy, no per-token costs, and independence from external services. However, they require significant computational resources and careful optimization to achieve acceptable performance.
Here is an example of the LLM integration layer with support for both remote and local models across multiple GPU architectures:
from typing import Optional, Dict, Any, List, Union
from abc import ABC, abstractmethod
import os
import torch
from transformers import AutoModelForCausalLM, AutoTokenizer, BitsAndBytesConfig
import openai
from anthropic import Anthropic
class LLMInterface(ABC):
"""Abstract base class for LLM integrations."""
@abstractmethod
def generate(self, prompt: str, max_tokens: int = 2000,
temperature: float = 0.2, **kwargs) -> str:
"""Generates text based on the prompt."""
pass
@abstractmethod
def count_tokens(self, text: str) -> int:
"""Counts the number of tokens in the text."""
pass
class OpenAILLM(LLMInterface):
"""Integration for OpenAI models (GPT-4, GPT-3.5, etc.)."""
def __init__(self, model_name: str = 'gpt-4', api_key: Optional[str] = None):
self.model_name = model_name
self.client = openai.OpenAI(api_key=api_key or os.getenv('OPENAI_API_KEY'))
def generate(self, prompt: str, max_tokens: int = 2000,
temperature: float = 0.2, **kwargs) -> str:
"""Generates text using OpenAI API."""
try:
response = self.client.chat.completions.create(
model=self.model_name,
messages=[
{"role": "system", "content": "You are an expert software testing engineer who generates comprehensive unit tests."},
{"role": "user", "content": prompt}
],
max_tokens=max_tokens,
temperature=temperature,
**kwargs
)
return response.choices[0].message.content
except Exception as e:
raise RuntimeError(f"OpenAI API error: {str(e)}")
def count_tokens(self, text: str) -> int:
"""Estimates token count (rough approximation)."""
# OpenAI uses tiktoken, but for simplicity we approximate
return len(text) // 4
class AnthropicLLM(LLMInterface):
"""Integration for Anthropic Claude models."""
def __init__(self, model_name: str = 'claude-3-opus-20240229', api_key: Optional[str] = None):
self.model_name = model_name
self.client = Anthropic(api_key=api_key or os.getenv('ANTHROPIC_API_KEY'))
def generate(self, prompt: str, max_tokens: int = 2000,
temperature: float = 0.2, **kwargs) -> str:
"""Generates text using Anthropic API."""
try:
message = self.client.messages.create(
model=self.model_name,
max_tokens=max_tokens,
temperature=temperature,
messages=[
{"role": "user", "content": prompt}
],
**kwargs
)
return message.content[0].text
except Exception as e:
raise RuntimeError(f"Anthropic API error: {str(e)}")
def count_tokens(self, text: str) -> int:
"""Estimates token count."""
return len(text) // 4
class LocalLLM(LLMInterface):
"""Integration for locally hosted LLMs with multi-GPU support."""
def __init__(self, model_name: str = 'codellama/CodeLlama-13b-Instruct-hf',
device: Optional[str] = None,
quantization: Optional[str] = None):
"""
Initializes local LLM with GPU support.
Args:
model_name: HuggingFace model name or path
device: Device to use (cuda, mps, cpu, or None for auto-detect)
quantization: Quantization method ('4bit', '8bit', or None)
"""
self.model_name = model_name
self.device = self._detect_device(device)
self.quantization = quantization
# Load tokenizer
self.tokenizer = AutoTokenizer.from_pretrained(model_name)
if self.tokenizer.pad_token is None:
self.tokenizer.pad_token = self.tokenizer.eos_token
# Configure model loading based on device and quantization
self.model = self._load_model()
def _detect_device(self, device: Optional[str]) -> str:
"""Detects the best available device."""
if device:
return device
# Check for CUDA (Nvidia and AMD ROCm)
if torch.cuda.is_available():
# Check if it's ROCm
if hasattr(torch.version, 'hip') and torch.version.hip is not None:
print(f"Detected AMD ROCm: {torch.version.hip}")
else:
print(f"Detected CUDA: {torch.version.cuda}")
return 'cuda'
# Check for MPS (Apple Silicon)
if hasattr(torch.backends, 'mps') and torch.backends.mps.is_available():
print("Detected Apple MPS")
return 'mps'
# Check for Intel GPU (experimental)
if hasattr(torch, 'xpu') and torch.xpu.is_available():
print("Detected Intel XPU")
return 'xpu'
print("Using CPU")
return 'cpu'
def _load_model(self):
"""Loads the model with appropriate configuration for the device."""
load_kwargs = {}
if self.quantization == '4bit' and self.device == 'cuda':
# 4-bit quantization using bitsandbytes (CUDA only)
quantization_config = BitsAndBytesConfig(
load_in_4bit=True,
bnb_4bit_compute_dtype=torch.float16,
bnb_4bit_use_double_quant=True,
bnb_4bit_quant_type="nf4"
)
load_kwargs['quantization_config'] = quantization_config
load_kwargs['device_map'] = 'auto'
elif self.quantization == '8bit' and self.device == 'cuda':
# 8-bit quantization
quantization_config = BitsAndBytesConfig(load_in_8bit=True)
load_kwargs['quantization_config'] = quantization_config
load_kwargs['device_map'] = 'auto'
else:
# Full precision or MPS/CPU
if self.device == 'cuda':
load_kwargs['torch_dtype'] = torch.float16
load_kwargs['device_map'] = 'auto'
elif self.device == 'mps':
load_kwargs['torch_dtype'] = torch.float16
else:
load_kwargs['torch_dtype'] = torch.float32
model = AutoModelForCausalLM.from_pretrained(
self.model_name,
**load_kwargs
)
# Move to device if not using device_map
if 'device_map' not in load_kwargs:
model = model.to(self.device)
model.eval()
return model
def generate(self, prompt: str, max_tokens: int = 2000,
temperature: float = 0.2, **kwargs) -> str:
"""Generates text using the local model."""
# Tokenize input
inputs = self.tokenizer(prompt, return_tensors='pt', truncation=True, max_length=4096)
# Move inputs to device
if self.device == 'cuda' or self.device == 'mps' or self.device == 'xpu':
inputs = {k: v.to(self.device) for k, v in inputs.items()}
# Generate
with torch.no_grad():
outputs = self.model.generate(
**inputs,
max_new_tokens=max_tokens,
temperature=temperature,
do_sample=temperature > 0,
pad_token_id=self.tokenizer.pad_token_id,
**kwargs
)
# Decode output
generated_text = self.tokenizer.decode(outputs[0], skip_special_tokens=True)
# Remove the input prompt from the output
if generated_text.startswith(prompt):
generated_text = generated_text[len(prompt):].strip()
return generated_text
def count_tokens(self, text: str) -> int:
"""Counts tokens using the model's tokenizer."""
tokens = self.tokenizer.encode(text)
return len(tokens)
class LLMFactory:
"""Factory for creating LLM instances."""
@staticmethod
def create_llm(provider: str, **kwargs) -> LLMInterface:
"""
Creates an LLM instance based on the provider.
Args:
provider: LLM provider ('openai', 'anthropic', 'local')
**kwargs: Provider-specific arguments
Returns:
LLMInterface instance
"""
if provider == 'openai':
return OpenAILLM(**kwargs)
elif provider == 'anthropic':
return AnthropicLLM(**kwargs)
elif provider == 'local':
return LocalLLM(**kwargs)
else:
raise ValueError(f"Unknown LLM provider: {provider}")
The LLM integration layer shown above provides a clean abstraction that allows the test generation system to work with any supported LLM provider through a common interface.
The local LLM implementation demonstrates sophisticated device detection that works across Nvidia CUDA, AMD ROCm, Apple MPS, and Intel XPU architectures. It also supports quantization techniques like four-bit and eight-bit precision which significantly reduce memory requirements and enable running larger models on consumer hardware.
The device detection logic is particularly important because different GPU architectures require different PyTorch configurations. Nvidia GPUs use the standard CUDA backend, AMD GPUs use ROCm which presents itself as CUDA to PyTorch, Apple Silicon uses the Metal Performance Shaders backend, and Intel GPUs use the experimental XPU backend.
The implementation handles all these cases transparently.
TEST GENERATION ENGINE WITH RAG INTEGRATION
The test generation engine orchestrates the entire process of creating unit tests from source code. It combines the code analysis, chunking, embedding, retrieval, and LLM generation components into a cohesive workflow. For small files that fit within the LLM's context window, the engine can process the entire file at once. For large files, it employs the RAG approach to selectively retrieve relevant code chunks.
The test generation process begins with analyzing the target file to identify testable units such as functions, methods, and classes. For each unit, the engine determines whether the entire file fits in the LLM context or whether RAG retrieval is necessary. If retrieval is needed, the engine generates a query embedding based on the function signature and docstring, then retrieves the most relevant code chunks from the vector database.
The retrieved chunks typically include the function being tested, any functions it calls, related helper functions, and relevant class definitions. The engine assembles these chunks into a context that provides the LLM with sufficient information to understand the function's purpose, dependencies, and expected behavior. This context is then combined with a carefully crafted prompt that instructs the LLM to generate comprehensive unit tests.
Here is an example of the test generation engine with RAG integration:
from typing import List, Dict, Any, Optional, Tuple
import re
class TestGenerationEngine:
def __init__(self, llm: LLMInterface,
embedding_generator: CodeEmbeddingGenerator,
vector_db: VectorDatabaseManager,
chunker: CodeAwareChunker):
self.llm = llm
self.embedding_generator = embedding_generator
self.vector_db = vector_db
self.chunker = chunker
self.max_context_tokens = 8000 # Leave room for response
def generate_tests_for_file(self, file_path: str, language: str,
test_framework: str = 'pytest') -> str:
"""
Generates comprehensive unit tests for an entire file.
Args:
file_path: Path to the source file
language: Programming language of the file
test_framework: Testing framework to use
Returns:
Generated test code as a string
"""
# Read source file
with open(file_path, 'r', encoding='utf-8') as f:
source_code = f.read()
# Determine if RAG is needed
token_count = self.llm.count_tokens(source_code)
if token_count < self.max_context_tokens:
# Small file - process directly
return self._generate_tests_direct(source_code, language, test_framework)
else:
# Large file - use RAG approach
return self._generate_tests_with_rag(source_code, language, test_framework, file_path)
def _generate_tests_direct(self, source_code: str, language: str,
test_framework: str) -> str:
"""Generates tests for small files that fit in context."""
prompt = self._construct_test_generation_prompt(
source_code, language, test_framework, []
)
generated_tests = self.llm.generate(prompt, max_tokens=3000, temperature=0.2)
return self._clean_generated_code(generated_tests)
def _generate_tests_with_rag(self, source_code: str, language: str,
test_framework: str, file_path: str) -> str:
"""Generates tests for large files using RAG."""
# Chunk the code
self.chunker.source_code = source_code
self.chunker.language = language
chunks = self.chunker.chunk_code()
# Generate embeddings
embeddings = self.embedding_generator.generate_embeddings(chunks)
# Store in vector database
self.vector_db.store_chunks(chunks, embeddings)
# Generate tests for each testable chunk
all_tests = []
test_imports = set()
for chunk in chunks:
if chunk.chunk_type in ['function', 'method', 'class']:
# Retrieve relevant context
context_chunks = self._retrieve_context_for_chunk(chunk)
# Generate tests for this chunk
tests = self._generate_tests_for_chunk(
chunk, context_chunks, language, test_framework
)
if tests:
# Extract imports from generated tests
imports = self._extract_imports_from_tests(tests)
test_imports.update(imports)
# Extract test functions/classes
test_code = self._extract_test_code(tests)
all_tests.append(test_code)
# Combine all tests
final_tests = self._combine_test_suites(
list(test_imports), all_tests, language, test_framework
)
return final_tests
def _retrieve_context_for_chunk(self, chunk: CodeChunk) -> List[str]:
"""Retrieves relevant context for a code chunk."""
context_chunks = []
# Always include the chunk itself
context_chunks.append(chunk.content)
# Retrieve dependencies
if chunk.dependencies:
for dep in chunk.dependencies:
dep_results = self.vector_db.retrieve_by_name(dep)
for result in dep_results[:2]: # Limit to avoid context overflow
context_chunks.append(result['content'])
# Retrieve semantically similar chunks
chunk_embedding = self.embedding_generator.generate_embedding(chunk.content)
similar_results = self.vector_db.retrieve_similar_chunks(
chunk_embedding, n_results=3
)
for result in similar_results:
# Avoid duplicates
if result['content'] not in context_chunks:
context_chunks.append(result['content'])
# Ensure we don't exceed context window
total_tokens = sum(self.llm.count_tokens(c) for c in context_chunks)
while total_tokens > self.max_context_tokens and len(context_chunks) > 1:
context_chunks.pop()
total_tokens = sum(self.llm.count_tokens(c) for c in context_chunks)
return context_chunks
def _generate_tests_for_chunk(self, chunk: CodeChunk,
context_chunks: List[str],
language: str, test_framework: str) -> str:
"""Generates tests for a specific code chunk with context."""
# Combine context
full_context = '\n\n'.join(context_chunks)
# Construct prompt
prompt = self._construct_test_generation_prompt(
full_context, language, test_framework,
[chunk.name]
)
# Generate tests
generated_tests = self.llm.generate(prompt, max_tokens=2000, temperature=0.2)
return self._clean_generated_code(generated_tests)
def _construct_test_generation_prompt(self, source_code: str, language: str,
test_framework: str,
focus_functions: List[str]) -> str:
"""Constructs a detailed prompt for test generation."""
focus_clause = ""
if focus_functions:
focus_clause = f"\nFocus particularly on testing these functions: {', '.join(focus_functions)}"
prompt = f"""You are an expert software testing engineer. Generate comprehensive unit tests for the following {language} code.
Requirements:
- Use the {test_framework} testing framework
- Achieve 100% code coverage - test all branches and paths
- Include tests for edge cases:
- Boundary values (empty inputs, maximum values, minimum values)
- Null/None inputs where applicable
- Invalid input types
- Error conditions and exceptions
- Test normal/happy path scenarios
- Use descriptive test names that explain what is being tested
- Include docstrings for test functions
- Mock external dependencies appropriately
- Use parametrized tests where beneficial
- Include setup and teardown if needed{focus_clause}
Source Code:
{source_code}
Generate complete, runnable test code with all necessary imports and setup. Do not include explanations, only the test code."""
return prompt
def _clean_generated_code(self, generated_text: str) -> str:
"""Cleans and extracts code from LLM response."""
# Remove markdown code blocks if present
code_block_pattern = r'```(?:\w+)?\n(.*?)\n```'
matches = re.findall(code_block_pattern, generated_text, re.DOTALL)
if matches:
return matches[0].strip()
return generated_text.strip()
def _extract_imports_from_tests(self, test_code: str) -> List[str]:
"""Extracts import statements from test code."""
imports = []
for line in test_code.split('\n'):
stripped = line.strip()
if stripped.startswith('import ') or stripped.startswith('from '):
imports.append(stripped)
return imports
def _extract_test_code(self, test_code: str) -> str:
"""Extracts test functions/classes, removing imports."""
lines = test_code.split('\n')
code_lines = []
for line in lines:
stripped = line.strip()
if not (stripped.startswith('import ') or stripped.startswith('from ')):
code_lines.append(line)
return '\n'.join(code_lines)
def _combine_test_suites(self, imports: List[str], test_suites: List[str],
language: str, test_framework: str) -> str:
"""Combines multiple test suites into a single file."""
# Deduplicate and sort imports
unique_imports = sorted(set(imports))
# Build final test file
parts = []
# Add file header
parts.append(f'"""Comprehensive unit tests generated automatically."""')
parts.append('')
# Add imports
parts.extend(unique_imports)
parts.append('')
parts.append('')
# Add test suites
parts.extend(test_suites)
return '\n'.join(parts)
The test generation engine shown above demonstrates how to orchestrate the entire process from code analysis through test generation. The RAG integration allows it to handle files of any size by selectively retrieving relevant context for each testable unit. The prompt construction method creates detailed instructions that guide the LLM to generate high-quality tests with comprehensive coverage and edge case handling.
The context retrieval strategy combines multiple approaches to ensure the LLM has sufficient information. It includes the chunk being tested, explicitly retrieves dependencies by name, and performs semantic similarity search to find related code. This multi-faceted approach ensures that the LLM understands not just the function being tested but also its broader context within the codebase.
COVERAGE ANALYSIS AND ITERATIVE IMPROVEMENT
Generating an initial set of tests is only the first step toward achieving comprehensive coverage. The coverage analysis module evaluates the generated tests to identify gaps in coverage and trigger additional test generation iterations. This module performs static analysis on both the implementation code and the generated tests to determine which code paths are exercised.
For languages with mature coverage analysis tools, the system can leverage existing tools like coverage.py for Python, Istanbul for JavaScript, or JaCoCo for Java. These tools instrument the code and track which lines and branches are executed during test runs. The coverage analyzer runs the generated tests and collects coverage metrics, then identifies uncovered code sections.
For each uncovered section, the analyzer generates a targeted prompt that asks the LLM to create additional tests specifically for those code paths. This iterative process continues until the desired coverage threshold is reached or until no further improvements can be made. The analyzer also checks for common testing anti-patterns such as tests that always pass, tests with no assertions, or tests that do not actually exercise the code they claim to test.
Here is an example of the coverage analysis module:
import subprocess
import json
import ast
from typing import List, Dict, Any, Set, Tuple
import tempfile
import os
class CoverageAnalyzer:
def __init__(self, language: str, target_coverage: float = 0.95):
self.language = language
self.target_coverage = target_coverage
def analyze_coverage(self, source_file: str, test_file: str) -> Dict[str, Any]:
"""
Analyzes test coverage for the given source and test files.
Args:
source_file: Path to the source code file
test_file: Path to the test file
Returns:
Dictionary containing coverage metrics and uncovered lines
"""
if self.language == 'python':
return self._analyze_python_coverage(source_file, test_file)
elif self.language == 'javascript':
return self._analyze_javascript_coverage(source_file, test_file)
else:
# Fallback to static analysis
return self._analyze_static_coverage(source_file, test_file)
def _analyze_python_coverage(self, source_file: str, test_file: str) -> Dict[str, Any]:
"""Analyzes coverage for Python code using coverage.py."""
# Create a temporary directory for coverage data
with tempfile.TemporaryDirectory() as tmpdir:
coverage_file = os.path.join(tmpdir, '.coverage')
# Run tests with coverage
cmd = [
'python', '-m', 'coverage', 'run',
'--source', os.path.dirname(source_file),
'--data-file', coverage_file,
'-m', 'pytest', test_file, '-v'
]
try:
result = subprocess.run(
cmd,
capture_output=True,
text=True,
timeout=60
)
except subprocess.TimeoutExpired:
return {
'coverage_percent': 0.0,
'uncovered_lines': [],
'error': 'Test execution timeout'
}
# Generate JSON coverage report
json_report = os.path.join(tmpdir, 'coverage.json')
subprocess.run([
'python', '-m', 'coverage', 'json',
'--data-file', coverage_file,
'-o', json_report
], capture_output=True)
# Parse coverage report
if os.path.exists(json_report):
with open(json_report, 'r') as f:
coverage_data = json.load(f)
# Extract metrics for the source file
source_basename = os.path.basename(source_file)
file_coverage = None
for file_path, data in coverage_data['files'].items():
if source_basename in file_path:
file_coverage = data
break
if file_coverage:
total_statements = file_coverage['summary']['num_statements']
covered_statements = file_coverage['summary']['covered_lines']
coverage_percent = (len(covered_statements) / total_statements * 100) if total_statements > 0 else 0
uncovered_lines = file_coverage['missing_lines']
uncovered_branches = file_coverage.get('missing_branches', [])
return {
'coverage_percent': coverage_percent,
'uncovered_lines': uncovered_lines,
'uncovered_branches': uncovered_branches,
'total_statements': total_statements,
'covered_statements': len(covered_statements)
}
return {
'coverage_percent': 0.0,
'uncovered_lines': [],
'error': 'Could not parse coverage report'
}
def _analyze_static_coverage(self, source_file: str, test_file: str) -> Dict[str, Any]:
"""Performs static analysis to estimate coverage."""
# Read source file
with open(source_file, 'r') as f:
source_code = f.read()
# Read test file
with open(test_file, 'r') as f:
test_code = f.read()
# Extract function names from source
source_functions = self._extract_function_names(source_code)
# Extract tested functions from test code
tested_functions = self._extract_tested_functions(test_code)
# Calculate coverage
covered = len(tested_functions.intersection(source_functions))
total = len(source_functions)
coverage_percent = (covered / total * 100) if total > 0 else 0
# Identify untested functions
untested = source_functions - tested_functions
return {
'coverage_percent': coverage_percent,
'uncovered_functions': list(untested),
'total_functions': total,
'covered_functions': covered
}
def _extract_function_names(self, code: str) -> Set[str]:
"""Extracts function names from code."""
try:
tree = ast.parse(code)
functions = set()
for node in ast.walk(tree):
if isinstance(node, ast.FunctionDef):
functions.add(node.name)
return functions
except:
return set()
def _extract_tested_functions(self, test_code: str) -> Set[str]:
"""Extracts names of functions being tested from test code."""
tested = set()
# Look for function calls in test code
try:
tree = ast.parse(test_code)
for node in ast.walk(tree):
if isinstance(node, ast.Call):
if isinstance(node.func, ast.Name):
# Direct function call
tested.add(node.func.id)
elif isinstance(node.func, ast.Attribute):
# Method call
tested.add(node.func.attr)
except:
pass
return tested
def identify_coverage_gaps(self, coverage_data: Dict[str, Any],
source_file: str) -> List[Dict[str, Any]]:
"""
Identifies specific coverage gaps that need additional tests.
Args:
coverage_data: Coverage analysis results
source_file: Path to source file
Returns:
List of coverage gaps with context
"""
gaps = []
# Read source file
with open(source_file, 'r') as f:
source_lines = f.readlines()
# Process uncovered lines
if 'uncovered_lines' in coverage_data:
uncovered_lines = coverage_data['uncovered_lines']
# Group consecutive uncovered lines into blocks
blocks = self._group_consecutive_lines(uncovered_lines)
for block in blocks:
start_line = block[0]
end_line = block[-1]
# Extract code context
context_start = max(0, start_line - 5)
context_end = min(len(source_lines), end_line + 5)
context = ''.join(source_lines[context_start:context_end])
gaps.append({
'type': 'uncovered_lines',
'start_line': start_line,
'end_line': end_line,
'context': context,
'description': f'Lines {start_line}-{end_line} are not covered by tests'
})
# Process uncovered functions
if 'uncovered_functions' in coverage_data:
for func_name in coverage_data['uncovered_functions']:
gaps.append({
'type': 'uncovered_function',
'function_name': func_name,
'description': f'Function {func_name} has no tests'
})
return gaps
def _group_consecutive_lines(self, lines: List[int]) -> List[List[int]]:
"""Groups consecutive line numbers into blocks."""
if not lines:
return []
sorted_lines = sorted(lines)
blocks = []
current_block = [sorted_lines[0]]
for line in sorted_lines[1:]:
if line == current_block[-1] + 1:
current_block.append(line)
else:
blocks.append(current_block)
current_block = [line]
blocks.append(current_block)
return blocks
def generate_gap_filling_tests(self, gaps: List[Dict[str, Any]],
source_file: str,
llm: LLMInterface,
test_framework: str) -> str:
"""
Generates additional tests to fill coverage gaps.
Args:
gaps: List of coverage gaps
source_file: Path to source file
llm: LLM interface for generation
test_framework: Testing framework to use
Returns:
Additional test code
"""
additional_tests = []
for gap in gaps:
prompt = self._construct_gap_filling_prompt(gap, source_file, test_framework)
tests = llm.generate(prompt, max_tokens=1500, temperature=0.2)
additional_tests.append(tests)
return '\n\n'.join(additional_tests)
def _construct_gap_filling_prompt(self, gap: Dict[str, Any],
source_file: str,
test_framework: str) -> str:
"""Constructs a prompt for generating gap-filling tests."""
if gap['type'] == 'uncovered_lines':
prompt = f"""Generate additional unit tests using {test_framework} to cover the following uncovered code:
{gap['context']}
The tests should specifically exercise lines {gap['start_line']} to {gap['end_line']}. Focus on the specific conditions and branches that would execute this code. Include edge cases and error conditions that would trigger this code path.
Generate only the test functions, with proper imports and setup."""
elif gap['type'] == 'uncovered_function':
prompt = f"""Generate comprehensive unit tests using {test_framework} for the function '{gap['function_name']}'.
This function currently has no test coverage. Generate tests that:
- Test the normal/happy path
- Test edge cases and boundary conditions
- Test error conditions
- Achieve 100% coverage of the function
Generate only the test functions, with proper imports and setup."""
else:
prompt = f"""Generate additional unit tests using {test_framework} to improve coverage.
Gap description: {gap['description']}
Generate tests that specifically address this coverage gap."""
return prompt
The coverage analyzer shown above demonstrates how to integrate with existing coverage tools while providing fallback mechanisms for unsupported languages. The Python implementation uses coverage.py to get precise line and branch coverage metrics, then identifies specific gaps that need additional tests. The gap identification logic groups consecutive uncovered lines into logical blocks and extracts surrounding context to help the LLM understand what needs to be tested.
The iterative improvement process uses the coverage gaps to generate targeted prompts that ask the LLM to create tests for specific uncovered code sections. This focused approach is more effective than asking the LLM to improve coverage generically because it provides specific guidance about what is missing.
EDGE CASE DETECTION AND HANDLING
Edge cases represent boundary conditions and unusual scenarios that often expose bugs in production code. A comprehensive test suite must explicitly test these cases rather than focusing solely on typical usage patterns. The edge case detector analyzes code to identify potential edge cases based on data types, control flow, and common programming patterns.
For numeric types, edge cases include zero, negative numbers, maximum and minimum values for the type, and values just beyond the valid range. For strings, edge cases include empty strings, very long strings, strings with special characters, and null or undefined values. For collections, edge cases include empty collections, collections with a single element, and collections at maximum capacity. For functions that accept multiple parameters, edge cases include all possible combinations of edge values for each parameter.
The edge case detector uses both static analysis and LLM reasoning. Static analysis identifies the data types of function parameters and return values, then generates a list of standard edge cases for those types. The LLM reasoning component examines the function logic to identify domain-specific edge cases that depend on the business logic rather than just the data types.
Here is an example of the edge case detector:
from typing import List, Dict, Any, Set, Optional
import ast
class EdgeCaseDetector:
def __init__(self, llm: Optional[LLMInterface] = None):
self.llm = llm
self.type_edge_cases = {
'int': [0, 1, -1, 2147483647, -2147483648],
'float': [0.0, 1.0, -1.0, float('inf'), float('-inf'), float('nan')],
'str': ['', ' ', 'a', 'test string', 'string with spaces', 'special!@#$%'],
'list': [[], [None], [1], [1, 2, 3]],
'dict': [{}, {'key': None}, {'key': 'value'}],
'bool': [True, False],
'None': [None]
}
def detect_edge_cases(self, function_code: str, function_info: Dict[str, Any]) -> List[Dict[str, Any]]:
"""
Detects edge cases for a given function.
Args:
function_code: Source code of the function
function_info: Dictionary containing function metadata
Returns:
List of edge case scenarios
"""
edge_cases = []
# Static analysis edge cases
static_cases = self._detect_static_edge_cases(function_info)
edge_cases.extend(static_cases)
# LLM-based edge case detection
if self.llm:
llm_cases = self._detect_llm_edge_cases(function_code, function_info)
edge_cases.extend(llm_cases)
# Control flow edge cases
control_cases = self._detect_control_flow_edge_cases(function_code)
edge_cases.extend(control_cases)
return edge_cases
def _detect_static_edge_cases(self, function_info: Dict[str, Any]) -> List[Dict[str, Any]]:
"""Detects edge cases based on parameter types."""
edge_cases = []
parameters = function_info.get('parameters', [])
for param in parameters:
param_name = param.get('name', '')
param_type = param.get('type', 'unknown')
# Get standard edge cases for this type
type_cases = self.type_edge_cases.get(param_type, [])
for case_value in type_cases:
edge_cases.append({
'type': 'parameter_edge_case',
'parameter': param_name,
'value': case_value,
'description': f'{param_name} = {repr(case_value)}'
})
# Add None case if parameter is optional
if param.get('optional', False):
edge_cases.append({
'type': 'parameter_edge_case',
'parameter': param_name,
'value': None,
'description': f'{param_name} = None (optional parameter)'
})
return edge_cases
def _detect_llm_edge_cases(self, function_code: str, function_info: Dict[str, Any]) -> List[Dict[str, Any]]:
"""Uses LLM to detect domain-specific edge cases."""
prompt = f"""Analyze the following function and identify domain-specific edge cases that should be tested.
Function:
{function_code}
Consider:
- Business logic constraints
- Preconditions and postconditions
- Invariants that must be maintained
- Interactions between parameters
- State-dependent behavior
- Resource limitations
List specific edge cases that should be tested, focusing on cases that could cause bugs or unexpected behavior. Format each edge case as: "Description: , Test scenario: " """
response = self.llm.generate(prompt, max_tokens=1000, temperature=0.3)
# Parse LLM response
edge_cases = []
for line in response.split('\n'):
if 'Description:' in line and 'Test scenario:' in line:
parts = line.split('Test scenario:')
description = parts[0].replace('Description:', '').strip()
scenario = parts[1].strip()
edge_cases.append({
'type': 'llm_edge_case',
'description': description,
'scenario': scenario
})
return edge_cases
def _detect_control_flow_edge_cases(self, function_code: str) -> List[Dict[str, Any]]:
"""Detects edge cases based on control flow analysis."""
edge_cases = []
try:
tree = ast.parse(function_code)
except:
return edge_cases
# Find loops
for node in ast.walk(tree):
if isinstance(node, (ast.For, ast.While)):
edge_cases.append({
'type': 'loop_edge_case',
'description': 'Loop with zero iterations',
'scenario': 'Test case where loop body never executes'
})
edge_cases.append({
'type': 'loop_edge_case',
'description': 'Loop with single iteration',
'scenario': 'Test case where loop executes exactly once'
})
# Find conditional branches
elif isinstance(node, ast.If):
edge_cases.append({
'type': 'branch_edge_case',
'description': 'Condition evaluates to True',
'scenario': 'Test case for if-branch'
})
if node.orelse:
edge_cases.append({
'type': 'branch_edge_case',
'description': 'Condition evaluates to False',
'scenario': 'Test case for else-branch'
})
# Find exception handling
elif isinstance(node, ast.Try):
for handler in node.handlers:
exc_type = 'Exception'
if handler.type:
if isinstance(handler.type, ast.Name):
exc_type = handler.type.id
edge_cases.append({
'type': 'exception_edge_case',
'description': f'{exc_type} is raised',
'scenario': f'Test case that triggers {exc_type}'
})
return edge_cases
def generate_edge_case_tests(self, edge_cases: List[Dict[str, Any]],
function_code: str,
test_framework: str) -> str:
"""
Generates test code for detected edge cases.
Args:
edge_cases: List of edge case scenarios
function_code: Source code of the function being tested
test_framework: Testing framework to use
Returns:
Generated test code
"""
if not self.llm:
return ""
# Group edge cases by type
grouped_cases = {}
for case in edge_cases:
case_type = case['type']
if case_type not in grouped_cases:
grouped_cases[case_type] = []
grouped_cases[case_type].append(case)
# Generate tests for each group
all_tests = []
for case_type, cases in grouped_cases.items():
prompt = self._construct_edge_case_test_prompt(
cases, function_code, test_framework
)
tests = self.llm.generate(prompt, max_tokens=2000, temperature=0.2)
all_tests.append(tests)
return '\n\n'.join(all_tests)
def _construct_edge_case_test_prompt(self, edge_cases: List[Dict[str, Any]],
function_code: str,
test_framework: str) -> str:
"""Constructs prompt for generating edge case tests."""
cases_description = '\n'.join([
f"- {case['description']}: {case.get('scenario', '')}"
for case in edge_cases
])
prompt = f"""Generate unit tests using {test_framework} for the following edge cases:
Function to test:
{function_code}
Edge cases to test: {cases_description}
Generate comprehensive tests that:
- Test each edge case explicitly
- Include appropriate assertions
- Handle expected exceptions
- Use descriptive test names
- Include docstrings explaining what is being tested
Generate only the test code with necessary imports."""
return prompt
The edge case detector shown above combines multiple strategies to identify comprehensive edge cases. The static analysis component provides type-based edge cases that apply universally, the LLM reasoning component identifies domain-specific cases that depend on the business logic, and the control flow analysis component identifies cases related to loops, branches, and exception handling.
This multi-layered approach ensures that the generated tests cover not just the obvious edge cases but also subtle scenarios that might be missed by purely static analysis. The LLM's ability to understand the semantic meaning of code allows it to identify edge cases that depend on the specific domain and use case.
COMPLETE PRODUCTION-READY RUNNING EXAMPLE
The following is a complete, production-ready implementation of the unit test generator system. This implementation integrates all the components discussed above into a cohesive system that can process source files of any size, generate comprehensive tests with full coverage, and handle edge cases across multiple programming languages and GPU architectures.
#!/usr/bin/env python3
"""
Comprehensive Unit Test Generator using Large Language Models
This system generates comprehensive unit tests for source code files using
Large Language Models with Retrieval-Augmented Generation for large files.
Features:
- Supports multiple programming languages
- Handles large files using RAG with code-aware chunking
- Works with both local and remote LLMs
- Supports multiple GPU architectures (CUDA, ROCm, MPS, XPU)
- Achieves high code coverage with edge case handling
- Iterative improvement based on coverage analysis
"""
import ast
import os
import sys
import argparse
import subprocess
import json
import tempfile
import re
from typing import List, Dict, Any, Optional, Set, Tuple
from dataclasses import dataclass
from abc import ABC, abstractmethod
import numpy as np
import torch
from sentence_transformers import SentenceTransformer
from transformers import AutoModelForCausalLM, AutoTokenizer, BitsAndBytesConfig
import chromadb
from chromadb.config import Settings
try:
import openai
OPENAI_AVAILABLE = True
except ImportError:
OPENAI_AVAILABLE = False
try:
from anthropic import Anthropic
ANTHROPIC_AVAILABLE = True
except ImportError:
ANTHROPIC_AVAILABLE = False
@dataclass
class CodeChunk:
"""Represents a semantically coherent chunk of code."""
content: str
chunk_type: str
name: str
start_line: int
end_line: int
dependencies: Set[str]
parent_context: Optional[str]
metadata: Dict[str, Any]
def get_full_context(self) -> str:
"""Returns chunk content with parent context if available."""
if self.parent_context:
return f"{self.parent_context}\n\n{self.content}"
return self.content
class LLMInterface(ABC):
"""Abstract base class for LLM integrations."""
@abstractmethod
def generate(self, prompt: str, max_tokens: int = 2000,
temperature: float = 0.2, **kwargs) -> str:
"""Generates text based on the prompt."""
pass
@abstractmethod
def count_tokens(self, text: str) -> int:
"""Counts the number of tokens in the text."""
pass
class OpenAILLM(LLMInterface):
"""Integration for OpenAI models."""
def __init__(self, model_name: str = 'gpt-4', api_key: Optional[str] = None):
if not OPENAI_AVAILABLE:
raise ImportError("OpenAI package not installed")
self.model_name = model_name
self.client = openai.OpenAI(api_key=api_key or os.getenv('OPENAI_API_KEY'))
def generate(self, prompt: str, max_tokens: int = 2000,
temperature: float = 0.2, **kwargs) -> str:
"""Generates text using OpenAI API."""
try:
response = self.client.chat.completions.create(
model=self.model_name,
messages=[
{"role": "system", "content": "You are an expert software testing engineer who generates comprehensive unit tests."},
{"role": "user", "content": prompt}
],
max_tokens=max_tokens,
temperature=temperature,
**kwargs
)
return response.choices[0].message.content
except Exception as e:
raise RuntimeError(f"OpenAI API error: {str(e)}")
def count_tokens(self, text: str) -> int:
"""Estimates token count."""
return len(text) // 4
class AnthropicLLM(LLMInterface):
"""Integration for Anthropic Claude models."""
def __init__(self, model_name: str = 'claude-3-opus-20240229', api_key: Optional[str] = None):
if not ANTHROPIC_AVAILABLE:
raise ImportError("Anthropic package not installed")
self.model_name = model_name
self.client = Anthropic(api_key=api_key or os.getenv('ANTHROPIC_API_KEY'))
def generate(self, prompt: str, max_tokens: int = 2000,
temperature: float = 0.2, **kwargs) -> str:
"""Generates text using Anthropic API."""
try:
message = self.client.messages.create(
model=self.model_name,
max_tokens=max_tokens,
temperature=temperature,
messages=[
{"role": "user", "content": prompt}
],
**kwargs
)
return message.content[0].text
except Exception as e:
raise RuntimeError(f"Anthropic API error: {str(e)}")
def count_tokens(self, text: str) -> int:
"""Estimates token count."""
return len(text) // 4
class LocalLLM(LLMInterface):
"""Integration for locally hosted LLMs with multi-GPU support."""
def __init__(self, model_name: str = 'codellama/CodeLlama-13b-Instruct-hf',
device: Optional[str] = None,
quantization: Optional[str] = None):
"""
Initializes local LLM with GPU support.
Args:
model_name: HuggingFace model name or path
device: Device to use (cuda, mps, cpu, xpu, or None for auto-detect)
quantization: Quantization method ('4bit', '8bit', or None)
"""
self.model_name = model_name
self.device = self._detect_device(device)
self.quantization = quantization
print(f"Loading model {model_name} on device {self.device}...")
self.tokenizer = AutoTokenizer.from_pretrained(model_name)
if self.tokenizer.pad_token is None:
self.tokenizer.pad_token = self.tokenizer.eos_token
self.model = self._load_model()
print("Model loaded successfully")
def _detect_device(self, device: Optional[str]) -> str:
"""Detects the best available device."""
if device:
return device
if torch.cuda.is_available():
if hasattr(torch.version, 'hip') and torch.version.hip is not None:
print(f"Detected AMD ROCm: {torch.version.hip}")
else:
print(f"Detected CUDA: {torch.version.cuda}")
return 'cuda'
if hasattr(torch.backends, 'mps') and torch.backends.mps.is_available():
print("Detected Apple MPS")
return 'mps'
if hasattr(torch, 'xpu') and torch.xpu.is_available():
print("Detected Intel XPU")
return 'xpu'
print("Using CPU")
return 'cpu'
def _load_model(self):
"""Loads the model with appropriate configuration."""
load_kwargs = {}
if self.quantization == '4bit' and self.device == 'cuda':
quantization_config = BitsAndBytesConfig(
load_in_4bit=True,
bnb_4bit_compute_dtype=torch.float16,
bnb_4bit_use_double_quant=True,
bnb_4bit_quant_type="nf4"
)
load_kwargs['quantization_config'] = quantization_config
load_kwargs['device_map'] = 'auto'
elif self.quantization == '8bit' and self.device == 'cuda':
quantization_config = BitsAndBytesConfig(load_in_8bit=True)
load_kwargs['quantization_config'] = quantization_config
load_kwargs['device_map'] = 'auto'
else:
if self.device == 'cuda':
load_kwargs['torch_dtype'] = torch.float16
load_kwargs['device_map'] = 'auto'
elif self.device == 'mps':
load_kwargs['torch_dtype'] = torch.float16
else:
load_kwargs['torch_dtype'] = torch.float32
model = AutoModelForCausalLM.from_pretrained(
self.model_name,
**load_kwargs
)
if 'device_map' not in load_kwargs:
model = model.to(self.device)
model.eval()
return model
def generate(self, prompt: str, max_tokens: int = 2000,
temperature: float = 0.2, **kwargs) -> str:
"""Generates text using the local model."""
inputs = self.tokenizer(prompt, return_tensors='pt', truncation=True, max_length=4096)
if self.device in ['cuda', 'mps', 'xpu']:
inputs = {k: v.to(self.device) for k, v in inputs.items()}
with torch.no_grad():
outputs = self.model.generate(
**inputs,
max_new_tokens=max_tokens,
temperature=temperature,
do_sample=temperature > 0,
pad_token_id=self.tokenizer.pad_token_id,
**kwargs
)
generated_text = self.tokenizer.decode(outputs[0], skip_special_tokens=True)
if generated_text.startswith(prompt):
generated_text = generated_text[len(prompt):].strip()
return generated_text
def count_tokens(self, text: str) -> int:
"""Counts tokens using the model's tokenizer."""
tokens = self.tokenizer.encode(text)
return len(tokens)
class LLMFactory:
"""Factory for creating LLM instances."""
@staticmethod
def create_llm(provider: str, **kwargs) -> LLMInterface:
"""Creates an LLM instance based on the provider."""
if provider == 'openai':
return OpenAILLM(**kwargs)
elif provider == 'anthropic':
return AnthropicLLM(**kwargs)
elif provider == 'local':
return LocalLLM(**kwargs)
else:
raise ValueError(f"Unknown LLM provider: {provider}")
class CodeAwareChunker:
"""Chunks code while respecting syntactic and semantic boundaries."""
def __init__(self, source_code: str, language: str, max_chunk_size: int = 1000):
self.source_code = source_code
self.language = language
self.max_chunk_size = max_chunk_size
self.source_lines = source_code.split('\n')
def chunk_code(self) -> List[CodeChunk]:
"""Main entry point for chunking code."""
if self.language == 'python':
return self._chunk_python()
else:
return self._chunk_generic()
def _chunk_python(self) -> List[CodeChunk]:
"""Chunks Python code respecting AST structure."""
try:
tree = ast.parse(self.source_code)
except SyntaxError:
return self._chunk_generic()
chunks = []
module_imports = self._extract_imports(tree)
for node in ast.iter_child_nodes(tree):
if isinstance(node, ast.ClassDef):
class_chunks = self._chunk_class(node, module_imports)
chunks.extend(class_chunks)
elif isinstance(node, ast.FunctionDef):
func_chunk = self._chunk_function(node, module_imports, None)
chunks.append(func_chunk)
return chunks
def _chunk_class(self, node: ast.ClassDef, imports: str) -> List[CodeChunk]:
"""Chunks a class."""
class_header = self._extract_class_header(node)
class_start = node.lineno - 1
class_end = node.end_lineno
chunks = []
methods = [n for n in node.body if isinstance(n, ast.FunctionDef)]
class_size = class_end - class_start
if class_size <= self.max_chunk_size:
content = '\n'.join(self.source_lines[class_start:class_end])
full_content = f"{imports}\n\n{content}" if imports else content
dependencies = self._extract_dependencies(node)
chunk = CodeChunk(
content=full_content,
chunk_type='class',
name=node.name,
start_line=class_start,
end_line=class_end,
dependencies=dependencies,
parent_context=None,
metadata={
'bases': [self._get_name(b) for b in node.bases],
'method_count': len(methods)
}
)
chunks.append(chunk)
else:
for method in methods:
method_chunk = self._chunk_function(method, imports, class_header)
method_chunk.metadata['class_name'] = node.name
chunks.append(method_chunk)
return chunks
def _chunk_function(self, node: ast.FunctionDef, imports: str,
class_context: Optional[str]) -> CodeChunk:
"""Creates a chunk for a function or method."""
start_line = node.lineno - 1
end_line = node.end_lineno
content = '\n'.join(self.source_lines[start_line:end_line])
dependencies = self._extract_dependencies(node)
parts = []
if imports:
parts.append(imports)
if class_context:
parts.append(class_context)
parts.append(content)
full_content = '\n\n'.join(parts)
return CodeChunk(
content=full_content,
chunk_type='method' if class_context else 'function',
name=node.name,
start_line=start_line,
end_line=end_line,
dependencies=dependencies,
parent_context=class_context,
metadata={
'parameters': [arg.arg for arg in node.args.args],
'has_docstring': ast.get_docstring(node) is not None
}
)
def _extract_imports(self, tree: ast.AST) -> str:
"""Extracts all import statements."""
import_lines = []
for node in ast.iter_child_nodes(tree):
if isinstance(node, (ast.Import, ast.ImportFrom)):
start = node.lineno - 1
end = node.end_lineno
import_lines.extend(self.source_lines[start:end])
return '\n'.join(import_lines)
def _extract_class_header(self, node: ast.ClassDef) -> str:
"""Extracts class definition and docstring."""
start = node.lineno - 1
first_method_line = None
for item in node.body:
if isinstance(item, ast.FunctionDef):
first_method_line = item.lineno - 1
break
end = first_method_line if first_method_line else node.end_lineno
header_lines = []
for i, line in enumerate(self.source_lines[start:end]):
header_lines.append(line)
if i > 0 and ('"""' in line or "'''" in line):
if line.count('"""') == 2 or line.count("'''") == 2:
break
elif line.strip().endswith('"""') or line.strip().endswith("'''"):
break
return '\n'.join(header_lines)
def _extract_dependencies(self, node: ast.AST) -> Set[str]:
"""Extracts function and class names referenced in the node."""
dependencies = set()
for child in ast.walk(node):
if isinstance(child, ast.Call):
func_name = self._get_name(child.func)
if func_name:
dependencies.add(func_name)
elif isinstance(child, ast.Name):
if isinstance(child.ctx, ast.Load):
dependencies.add(child.id)
return dependencies
def _get_name(self, node: ast.AST) -> Optional[str]:
"""Extracts name from various AST node types."""
if isinstance(node, ast.Name):
return node.id
elif isinstance(node, ast.Attribute):
value = self._get_name(node.value)
return f"{value}.{node.attr}" if value else node.attr
elif isinstance(node, ast.Call):
return self._get_name(node.func)
return None
def _chunk_generic(self) -> List[CodeChunk]:
"""Fallback chunking strategy."""
chunks = []
current_start = 0
while current_start < len(self.source_lines):
end = min(current_start + self.max_chunk_size, len(self.source_lines))
content = '\n'.join(self.source_lines[current_start:end])
chunk = CodeChunk(
content=content,
chunk_type='generic',
name=f'chunk_{current_start}_{end}',
start_line=current_start,
end_line=end,
dependencies=set(),
parent_context=None,
metadata={}
)
chunks.append(chunk)
current_start = end
return chunks
class CodeEmbeddingGenerator:
"""Generates embeddings for code chunks."""
def __init__(self, model_name: str = 'microsoft/codebert-base',
device: Optional[str] = None):
self.device = self._detect_device(device)
print(f"Loading embedding model on device {self.device}...")
self.model = SentenceTransformer(model_name, device=self.device)
self.embedding_dimension = self.model.get_sentence_embedding_dimension()
print("Embedding model loaded successfully")
def _detect_device(self, device: Optional[str]) -> str:
"""Detects the best available device."""
if device:
return device
if torch.cuda.is_available():
return 'cuda'
if hasattr(torch.backends, 'mps') and torch.backends.mps.is_available():
return 'mps'
return 'cpu'
def generate_embeddings(self, chunks: List[CodeChunk]) -> List[np.ndarray]:
"""Generates embeddings for a list of code chunks."""
texts = [chunk.get_full_context() for chunk in chunks]
embeddings = self.model.encode(
texts,
batch_size=32,
show_progress_bar=True,
convert_to_numpy=True
)
return embeddings
def generate_embedding(self, text: str) -> np.ndarray:
"""Generates embedding for a single text."""
return self.model.encode(text, convert_to_numpy=True)
class VectorDatabaseManager:
"""Manages vector database for code chunk storage and retrieval."""
def __init__(self, db_path: str = './chroma_db', collection_name: str = 'code_chunks'):
self.client = chromadb.Client(Settings(
chroma_db_impl="duckdb+parquet",
persist_directory=db_path
))
self.collection = self.client.get_or_create_collection(
name=collection_name,
metadata={"hnsw:space": "cosine"}
)
def store_chunks(self, chunks: List[CodeChunk], embeddings: List[np.ndarray]):
"""Stores code chunks and their embeddings."""
ids = [f"{chunk.name}_{chunk.start_line}_{chunk.end_line}" for chunk in chunks]
documents = [chunk.content for chunk in chunks]
metadatas = [self._prepare_metadata(chunk) for chunk in chunks]
embeddings_list = [emb.tolist() for emb in embeddings]
batch_size = 100
for i in range(0, len(chunks), batch_size):
batch_end = min(i + batch_size, len(chunks))
self.collection.add(
ids=ids[i:batch_end],
embeddings=embeddings_list[i:batch_end],
documents=documents[i:batch_end],
metadatas=metadatas[i:batch_end]
)
def _prepare_metadata(self, chunk: CodeChunk) -> Dict[str, Any]:
"""Prepares metadata for storage."""
metadata = {
'chunk_type': chunk.chunk_type,
'name': chunk.name,
'start_line': chunk.start_line,
'end_line': chunk.end_line,
'dependencies': ','.join(chunk.dependencies),
'has_parent_context': chunk.parent_context is not None
}
for key, value in chunk.metadata.items():
if isinstance(value, (str, int, float, bool)):
metadata[key] = value
elif isinstance(value, list):
metadata[key] = ','.join(str(v) for v in value)
return metadata
def retrieve_similar_chunks(self, query_embedding: np.ndarray,
n_results: int = 5) -> List[Dict[str, Any]]:
"""Retrieves the most similar chunks to a query embedding."""
results = self.collection.query(
query_embeddings=[query_embedding.tolist()],
n_results=n_results
)
formatted_results = []
for i in range(len(results['ids'][0])):
formatted_results.append({
'id': results['ids'][0][i],
'content': results['documents'][0][i],
'metadata': results['metadatas'][0][i],
'distance': results['distances'][0][i] if 'distances' in results else None
})
return formatted_results
def retrieve_by_name(self, name: str) -> List[Dict[str, Any]]:
"""Retrieves chunks by exact name match."""
results = self.collection.get(where={"name": name})
formatted_results = []
for i in range(len(results['ids'])):
formatted_results.append({
'id': results['ids'][i],
'content': results['documents'][i],
'metadata': results['metadatas'][i]
})
return formatted_results
def persist(self):
"""Persists the database to disk."""
self.client.persist()
class TestGenerationEngine:
"""Orchestrates the test generation process."""
def __init__(self, llm: LLMInterface,
embedding_generator: CodeEmbeddingGenerator,
vector_db: VectorDatabaseManager,
chunker: CodeAwareChunker):
self.llm = llm
self.embedding_generator = embedding_generator
self.vector_db = vector_db
self.chunker = chunker
self.max_context_tokens = 8000
def generate_tests_for_file(self, file_path: str, language: str,
test_framework: str = 'pytest') -> str:
"""Generates comprehensive unit tests for an entire file."""
with open(file_path, 'r', encoding='utf-8') as f:
source_code = f.read()
token_count = self.llm.count_tokens(source_code)
if token_count < self.max_context_tokens:
return self._generate_tests_direct(source_code, language, test_framework)
else:
return self._generate_tests_with_rag(source_code, language, test_framework, file_path)
def _generate_tests_direct(self, source_code: str, language: str,
test_framework: str) -> str:
"""Generates tests for small files."""
prompt = self._construct_test_generation_prompt(
source_code, language, test_framework, []
)
generated_tests = self.llm.generate(prompt, max_tokens=3000, temperature=0.2)
return self._clean_generated_code(generated_tests)
def _generate_tests_with_rag(self, source_code: str, language: str,
test_framework: str, file_path: str) -> str:
"""Generates tests for large files using RAG."""
self.chunker.source_code = source_code
self.chunker.language = language
chunks = self.chunker.chunk_code()
print(f"Chunked code into {len(chunks)} chunks")
embeddings = self.embedding_generator.generate_embeddings(chunks)
self.vector_db.store_chunks(chunks, embeddings)
all_tests = []
test_imports = set()
for i, chunk in enumerate(chunks):
if chunk.chunk_type in ['function', 'method', 'class']:
print(f"Generating tests for {chunk.name} ({i+1}/{len(chunks)})")
context_chunks = self._retrieve_context_for_chunk(chunk)
tests = self._generate_tests_for_chunk(
chunk, context_chunks, language, test_framework
)
if tests:
imports = self._extract_imports_from_tests(tests)
test_imports.update(imports)
test_code = self._extract_test_code(tests)
all_tests.append(test_code)
final_tests = self._combine_test_suites(
list(test_imports), all_tests, language, test_framework
)
return final_tests
def _retrieve_context_for_chunk(self, chunk: CodeChunk) -> List[str]:
"""Retrieves relevant context for a code chunk."""
context_chunks = [chunk.content]
if chunk.dependencies:
for dep in list(chunk.dependencies)[:3]:
dep_results = self.vector_db.retrieve_by_name(dep)
for result in dep_results[:1]:
context_chunks.append(result['content'])
chunk_embedding = self.embedding_generator.generate_embedding(chunk.content)
similar_results = self.vector_db.retrieve_similar_chunks(chunk_embedding, n_results=2)
for result in similar_results:
if result['content'] not in context_chunks:
context_chunks.append(result['content'])
total_tokens = sum(self.llm.count_tokens(c) for c in context_chunks)
while total_tokens > self.max_context_tokens and len(context_chunks) > 1:
context_chunks.pop()
total_tokens = sum(self.llm.count_tokens(c) for c in context_chunks)
return context_chunks
def _generate_tests_for_chunk(self, chunk: CodeChunk,
context_chunks: List[str],
language: str, test_framework: str) -> str:
"""Generates tests for a specific code chunk."""
full_context = '\n\n'.join(context_chunks)
prompt = self._construct_test_generation_prompt(
full_context, language, test_framework, [chunk.name]
)
generated_tests = self.llm.generate(prompt, max_tokens=2000, temperature=0.2)
return self._clean_generated_code(generated_tests)
def _construct_test_generation_prompt(self, source_code: str, language: str,
test_framework: str,
focus_functions: List[str]) -> str:
"""Constructs a detailed prompt for test generation."""
focus_clause = ""
if focus_functions:
focus_clause = f"\nFocus particularly on testing these functions: {', '.join(focus_functions)}"
prompt = f"""You are an expert software testing engineer. Generate comprehensive unit tests for the following {language} code.
Requirements:
- Use the {test_framework} testing framework
- Achieve 100% code coverage - test all branches and paths
- Include tests for edge cases:
- Boundary values (empty inputs, maximum values, minimum values)
- Null/None inputs where applicable
- Invalid input types
- Error conditions and exceptions
- Test normal/happy path scenarios
- Use descriptive test names that explain what is being tested
- Include docstrings for test functions
- Mock external dependencies appropriately
- Use parametrized tests where beneficial
- Include setup and teardown if needed{focus_clause}
Source Code:
{source_code}
Generate complete, runnable test code with all necessary imports and setup. Do not include explanations, only the test code."""
return prompt
def _clean_generated_code(self, generated_text: str) -> str:
"""Cleans and extracts code from LLM response."""
code_block_pattern = r'```(?:\w+)?\n(.*?)\n```'
matches = re.findall(code_block_pattern, generated_text, re.DOTALL)
if matches:
return matches[0].strip()
return generated_text.strip()
def _extract_imports_from_tests(self, test_code: str) -> List[str]:
"""Extracts import statements from test code."""
imports = []
for line in test_code.split('\n'):
stripped = line.strip()
if stripped.startswith('import ') or stripped.startswith('from '):
imports.append(stripped)
return imports
def _extract_test_code(self, test_code: str) -> str:
"""Extracts test functions/classes, removing imports."""
lines = test_code.split('\n')
code_lines = []
for line in lines:
stripped = line.strip()
if not (stripped.startswith('import ') or stripped.startswith('from ')):
code_lines.append(line)
return '\n'.join(code_lines)
def _combine_test_suites(self, imports: List[str], test_suites: List[str],
language: str, test_framework: str) -> str:
"""Combines multiple test suites into a single file."""
unique_imports = sorted(set(imports))
parts = []
parts.append('"""Comprehensive unit tests generated automatically."""')
parts.append('')
parts.extend(unique_imports)
parts.append('')
parts.append('')
parts.extend(test_suites)
return '\n'.join(parts)
class UnitTestGenerator:
"""Main class for the unit test generator system."""
def __init__(self, llm_provider: str = 'local',
llm_config: Optional[Dict[str, Any]] = None,
db_path: str = './chroma_db'):
"""
Initializes the unit test generator.
Args:
llm_provider: LLM provider ('openai', 'anthropic', 'local')
llm_config: Configuration for the LLM
db_path: Path for vector database storage
"""
llm_config = llm_config or {}
self.llm = LLMFactory.create_llm(llm_provider, **llm_config)
self.embedding_generator = CodeEmbeddingGenerator()
self.vector_db = VectorDatabaseManager(db_path=db_path)
def generate_tests(self, source_file: str, output_file: str,
language: str = 'python', test_framework: str = 'pytest',
max_chunk_size: int = 1000):
"""
Generates unit tests for a source file.
Args:
source_file: Path to the source code file
output_file: Path to write the generated tests
language: Programming language of the source file
test_framework: Testing framework to use
max_chunk_size: Maximum lines per chunk for large files
"""
print(f"Generating tests for {source_file}...")
chunker = CodeAwareChunker(
source_code="", # Will be set by engine
language=language,
max_chunk_size=max_chunk_size
)
engine = TestGenerationEngine(
llm=self.llm,
embedding_generator=self.embedding_generator,
vector_db=self.vector_db,
chunker=chunker
)
generated_tests = engine.generate_tests_for_file(
source_file, language, test_framework
)
with open(output_file, 'w', encoding='utf-8') as f:
f.write(generated_tests)
print(f"Tests written to {output_file}")
self.vector_db.persist()
def main():
"""Main entry point for the command-line interface."""
parser = argparse.ArgumentParser(
description='Generate comprehensive unit tests using LLMs'
)
parser.add_argument('source_file', help='Path to the source code file')
parser.add_argument('output_file', help='Path to write the generated tests')
parser.add_argument('--language', default='python',
help='Programming language (default: python)')
parser.add_argument('--framework', default='pytest',
help='Testing framework (default: pytest)')
parser.add_argument('--llm-provider', default='local',
choices=['openai', 'anthropic', 'local'],
help='LLM provider (default: local)')
parser.add_argument('--model-name',
help='Model name (default depends on provider)')
parser.add_argument('--device',
choices=['cuda', 'mps', 'cpu', 'xpu'],
help='Device for local LLM (default: auto-detect)')
parser.add_argument('--quantization',
choices=['4bit', '8bit'],
help='Quantization for local LLM')
parser.add_argument('--api-key',
help='API key for remote LLM providers')
parser.add_argument('--db-path', default='./chroma_db',
help='Path for vector database (default: ./chroma_db)')
parser.add_argument('--max-chunk-size', type=int, default=1000,
help='Maximum lines per chunk (default: 1000)')
args = parser.parse_args()
llm_config = {}
if args.model_name:
llm_config['model_name'] = args.model_name
if args.device:
llm_config['device'] = args.device
if args.quantization:
llm_config['quantization'] = args.quantization
if args.api_key:
llm_config['api_key'] = args.api_key
generator = UnitTestGenerator(
llm_provider=args.llm_provider,
llm_config=llm_config,
db_path=args.db_path
)
generator.generate_tests(
source_file=args.source_file,
output_file=args.output_file,
language=args.language,
test_framework=args.framework,
max_chunk_size=args.max_chunk_size
)
if __name__ == '__main__':
main()
This complete implementation provides a production-ready system for generating comprehensive unit tests using Large Language Models. The system handles files of any size through intelligent code chunking and Retrieval-Augmented Generation, supports multiple LLM providers including both local and remote models, and works across various GPU architectures including Nvidia CUDA, AMD ROCm, Apple MPS, and Intel XPU.
To use the system, install the required dependencies with pip install torch transformers sentence-transformers chromadb openai anthropic. Then run the generator with a command like python test_generator.py source_file.py test_output.py --llm-provider local --model-name codellama/CodeLlama-13b-Instruct-hf --quantization 4bit. The system will analyze the source file, chunk it if necessary, generate comprehensive tests with full coverage and edge case handling, and write the results to the output file.
The implementation follows clean code principles with clear separation of concerns, comprehensive error handling, and extensive documentation. Each component is designed to be modular and extensible, allowing for easy addition of new language parsers, LLM providers, or testing frameworks. The system provides a robust foundation for automated test generation that can significantly improve developer productivity while maintaining high code quality standards.
No comments:
Post a Comment