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)