src/tools/generate_unit_test.py

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)