Spaces:
Sleeping
Sleeping
File size: 4,835 Bytes
448c329 f3f56ef 808ecdb 448c329 f3f56ef 448c329 f3f56ef 2dc77a4 448c329 f3f56ef 448c329 808ecdb 448c329 f3f56ef 448c329 f3f56ef 448c329 f3f56ef 448c329 f3f56ef 448c329 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 | import sys
import os
src_path = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
sys.path.append(f"{src_path}/tools")
from utilities import replace_markdown_section_including_markers
from anthropic import Anthropic
from dotenv import load_dotenv
import os
load_dotenv()
def generate_unit_test(code: str, source_code_path: str, test_root_dir_path: str,
additional_system_prompt: str = "", anthropic_api_key: str = "") -> str:
"""
Generates a Jest unit test file for the provided source code using the Anthropic API.
This function leverages the 'QA-Sentinel' automated test engineering bot to create
comprehensive unit tests aiming for 100% statement, branch, function, and line coverage.
Args:
code (str): The source code string for which to generate unit tests.
source_code_path (str): The full path to the original source code file. This is used
to provide context to the AI for accurate test generation,
especially for imports or file-system related operations.
test_root_dir_path (str): The root directory where the generated Jest test file
should be placed. The output file will be named
`[original_filename].Spec.ts` within this directory.
additional_system_prompt (str, optional): An optional string to append to the
default system prompt for the Anthropic API.
Use this to provide extra instructions or
context to the AI for test generation.
Defaults to an empty string.
anthropic_api_key (str, optional): An optional string to let guest use their own Anthropic key.
Returns:
str: The content of the generated Jest test file as a string, or an error message
if the code is not provided or an exception occurs during generation.
"""
if not code:
return "Please provide code to analyze."
try:
envAnthrophicApiKey = os.environ["ANTHROPIC_API_KEY"]
if envAnthrophicApiKey == "" and anthropic_api_key == "":
return "Please enter your Anthropic API key or set ANTHROPIC_API_KEY in your environment."
if anthropic_api_key != "":
envAnthrophicApiKey = anthropic_api_key
client = Anthropic(api_key=envAnthrophicApiKey)
system_prompt = """
You are 'QA-Sentinel', an automated test engineering bot. Your single-minded purpose is to ensure absolute code quality by achieving 100% statement, branch, function, and line coverage. Failure is not an option.
"""
if additional_system_prompt:
system_prompt += f"""
{additional_system_prompt}
"""
else:
system_prompt += "The test file should be named [original_filename].spec.ts."
print(system_prompt)
print("\n\n\n")
user_prompt = replace_markdown_section_including_markers(f'{src_path}/prompts/generate_unit_test_v1.md',
'[REPLACE_SOURCE_CODE_START]',
'[REPLACE_SOURCE_CODE_END]', code)
user_prompt += f"""
**Additional Context**:
- The source code is located in `{source_code_path}`
- The output Jest file must be placed in `{test_root_dir_path}`
"""
print(user_prompt)
print("\n\n\n")
resp = client.messages.create(
model="claude-3-7-sonnet-20250219",
messages=[{"role": "user", "content": user_prompt}],
system=system_prompt,
max_tokens=4000,
temperature=0,
)
output = resp.content[0].text
return output
except Exception as exc:
return f"Error during unit test generation: {exc}"
# For testing/debugging
if __name__ == "__main__":
sample_code = """
/**
* Applies a discount to a price.
* @param price The original price.
* @param discountRate The discount rate (e.g., 0.15 for 15% off).
* @returns The discounted price.
*/
function calculateDiscount(price: number, discountRate: number): number {
// No discount if rate is invalid (e.g., greater than 1, implying more than 100% discount)
if (discountRate > 1) {
return price;
}
return price * (1 - discountRate);
}
"""
sample_source_code_path = "src/calculate_discount.ts"
sample_test_root_dir_path = "test/"
print("Generating Unit Test from sample code...")
unit_test = generate_unit_test(sample_code, sample_source_code_path, sample_test_root_dir_path)
print("\n" + unit_test)
|