baseline_vanilla.py

import json
import os
import argparse
import re
from tqdm import tqdm
import subprocess

from test_case_runner import TestCaseRunner
from configs import Configs
from dataset import Dataset
from agents import Agent


class VanillaTestGeneralizer(Agent):
    def __init__(self, configs):
        super(VanillaTestGeneralizer, self).__init__(configs.llm_name, configs.openai_key, configs.openai_base_url)
        self.configs = configs
        self.max_retry = 5
        self.retry_count = 0

        self.gen_prefix = '```\npackage '
        self.gen_suffix = '```'
        self.system_prompt = f"""You may have memorized information from the GitHub repository '{configs.project_name}' (URL is {configs.project_url}). For this task, you must not use any of that memorized information in your responses. Instead, base your answers exclusively on the context I provide in the document. If your response would otherwise rely on memorized '{configs.project_name}' data, replace that content with generic or random information unrelated to '{configs.project_name}'."""

        self.test_runner = TestCaseRunner(configs, configs.test_case_run_log_dir)

    def generalize(self, target_focal_method, target_test_case_path, initial_test, gen_test_class_name, junit_version):
        gen_test_cases = self._generalize(target_focal_method, initial_test, gen_test_class_name, junit_version)
        gen_test_cases_results = []

        for each_gen_test_case in gen_test_cases:
            error_msg, test_status, final_test_log = self.run_test_case(each_gen_test_case.strip(), target_test_case_path)
            gen_test_cases_results.append((each_gen_test_case.strip(), test_status, final_test_log))

        return gen_test_cases_results

    def _generalize(self, target_focal_method, initial_test, gen_test_class_name, junit_version):
        prompt = f"""# Instruction\nGiven the Target Focal Method and its Initial Test Case, generate sufficient more test cases for it without any commentary.\n\n# Target Focal Method\n```\n{target_focal_method}\n```\n\n# Initial Test Case\n```\n{initial_test}\n```\n\n# Requirements\n- Each generated test case must be in Java language and use JUnit {junit_version} framework.\n- Each generated test must be self-contained, i.e., it should include all necessary imports, and contains only one test method annotated @Test.\n- Each generated test must begin with the exact prefix: "{self.gen_prefix}" and end with the exact suffix: "{self.gen_suffix}". Ensure that no additional text appears before the prefix or after the suffix.\n- Each generated test case must be compilable, runnable, and pass successfully, and its class name must be {gen_test_class_name}."""

        messages = [{'role': 'user', 'content': prompt}]        
        raw_response = self.get_response(messages)

        extracted_codes = self.extract_code_from_response(raw_response)
        if not extracted_codes:
            print(f"[WARNING] Failed to extract code from response:\n{raw_response}\n")
            if self.retry_count < self.max_retry:
                self.retry_count += 1
                print(f"[INFO] Regenerate... (Attempt {self.retry_count}/{self.max_retry})")
                return self._generalize(target_focal_method, initial_test, gen_test_class_name, junit_version)
            else:
                print(f"[WARNING] Reached max retry limit. Returning empty code.")
                self.retry_count = 0
                return []
        self.retry_count = 0
        return extracted_codes
    
    def extract_code_from_response(self, response: str):
        codes = re.findall(r'```java(.*?)```', response, re.DOTALL)
        if len(codes) == 0:
            codes = re.findall(r'```(.*?)```', response, re.DOTALL)
            if len(codes) == 0:
                print(f"[Warning] The response does not contain any code: {response}")
                return []

        return codes
    
    def run_test_case(self, test_case, test_case_path):
        def _extract_error_msg(log):
            error_msg = []
            stop_flag = False
            for each_line in log.split('\n'):
                if each_line.strip().startswith('[INFO]'):
                    continue
                if each_line.strip().startswith('[main]'):
                    continue
                if each_line.strip().startswith('[WARNING]'):
                    continue
                
                if each_line.strip().startswith('[ERROR] Tests run:'):
                    if stop_flag:
                        break
                    else:
                        stop_flag = True
                
                if each_line.strip().startswith('[ERROR] To see the full stack trace'):
                    break

                error_msg.append(each_line)

            error_msg = '\n'.join(error_msg)
            return error_msg

        compile_log, test_log, compile_success, execute_success = self.test_runner.compile_and_execute_test_case(test_case, test_case_path) 

        if not compile_success:
            error_msg = _extract_error_msg(compile_log)
            test_status = 'fail_compile'
        elif not execute_success:
            error_msg = _extract_error_msg(test_log)
            test_status = 'fail_execute'

            test_run_info = re.search(r'Tests run: (\d+), Failures: (\d+), Errors: (\d+), Skipped: (\d+)', test_log)
            if test_run_info is not None:
                test_run_info = test_run_info.groups()

                if int(test_run_info[0]) > 1:
                    print(f'[INFO] Multiple test methods in a single test case: {test_case_path}')

                success = int(test_run_info[0]) - int(test_run_info[1]) - int(test_run_info[2]) - int(test_run_info[3])
                if success > 0:
                    test_status = 'success'
                    error_msg = ""
                elif int(test_run_info[1]) > 0:
                    test_status = 'fail_pass'
                else:
                    test_status = 'fail_execute'

        else:
            error_msg = ""
            test_status = 'success'

        if test_status == 'success':
            final_test_log = ''
        elif test_status == 'fail_compile':
            final_test_log = compile_log
        else:
            final_test_log = test_log

        return error_msg, test_status, final_test_log


def generation_pipeline():
    # prepare the datasets
    dataset = Dataset(configs)
    print('Loading datasets...')
    coverage_data = dataset.load_coverage_data_jacoco()

    with open(configs.dataset_with_generalization_path, 'r') as f:
        scenario_data = json.load(f)

    fm_name2path_tc_path = dict()
    for each_coverage_sample in coverage_data:
        focal_method = each_coverage_sample.focal_method.strip()
        focal_method_name = each_coverage_sample.focal_method_name
        focal_method_path = each_coverage_sample.focal_file_path
        tc_path = each_coverage_sample.test_case_path
        old_project_dir = tc_path.split(f'/{configs.project_name}/')[0]
        tc_path = tc_path.replace(f'{old_project_dir}/{configs.project_name}', configs.project_path_no_test_file)
        fm_name2path_tc_path[(focal_method, focal_method_name)] = (focal_method_path, tc_path)

    # clean test cases in repos_removing_test folder
    test_case_paths = [initial_fm_info.test_case_path for initial_fm_info in coverage_data]
    clean_test_cases(test_case_paths)

    # prepare test generalizer
    test_generalizer = VanillaTestGeneralizer(configs)

    # prepare the save path
    gen_save_path = prepare_save_path(configs.generated_test_case_save_path)

    # start generating test case
    total_generated_test_cases = []

    for iter_idx, each_scenario in tqdm(enumerate(scenario_data), total=len(scenario_data), ncols=80, desc='Generating test cases'):
        scenario_idx = each_scenario['target_coverage_idx']
        if scenario_idx >= args.stop_at:
            print(f'[INFO] Early stop at {scenario_idx}\n')
            break

        if scenario_idx < args.start_at:
            continue
            
        if args.specify_ids and scenario_idx not in args.specify_ids:
            continue
        
        print(f'\n\n[INFO] Coverage index: {scenario_idx}\n\n')

        if 'same_focal_method' not in each_scenario['generalizations']:
            print(f'[INFO] No same_focal_method generalizations for target pair {scenario_idx}. Skip.\n')
            each_scenario['generated_test_cases'] = []
            
            total_generated_test_cases.append(each_scenario)
            with open(gen_save_path, 'w') as f:
                json.dump(total_generated_test_cases, f, indent=4)

            continue
        
        initial_fm_info = coverage_data[scenario_idx]

        project_name = initial_fm_info.project_name
        init_focal_method_name = initial_fm_info.focal_method_name
        init_test_case = initial_fm_info.test_case
        init_test_case_name = initial_fm_info.test_case_name
        gen_test_class_name = initial_fm_info.test_case_path.split('/')[-1].replace('.java', '')

        assert init_focal_method_name == each_scenario['focal_method_name'], f'Inconsistent focal_method_name: {init_focal_method_name} vs {each_scenario["focal_method_name"]}'
        assert init_test_case_name in each_scenario['target_test_case_name'], f'Inconsistent test_case_name: {each_scenario["target_test_case_name"]} vs {init_test_case_name}'

        for gen_type, _ in each_scenario['generalizations'].items():
            if gen_type != 'same_focal_method':
                continue
            
            test_cases_and_results = test_generalizer.generalize(
                target_focal_method=initial_fm_info.focal_method,
                target_test_case_path=initial_fm_info.test_case_path,
                initial_test=init_test_case,
                gen_test_class_name=gen_test_class_name,
                junit_version=args.junit_version
            )

            for gen_idx, each_result in enumerate(test_cases_and_results):
                each_test_case, test_status, test_log = each_result
                # save results
                original_path = initial_fm_info.test_case_path
                init_focal_method_name_pure = init_focal_method_name.split('::::')[1].split('(')[0]
                generated_path = original_path[:-5] + f'#{init_focal_method_name_pure}#{init_test_case_name}{gen_idx}.java'
                total_generated_test_cases.append({
                    'project_name': project_name,
                    'target_coverage_idx': scenario_idx,
                    'focal_file_path': initial_fm_info.focal_file_path,
                    'focal_method_name': init_focal_method_name,
                    'generated_test_case': each_test_case,
                    'target_coverage': ''.join(each_scenario['target_coverage']),
                    'target_context': each_scenario['target_context'],
                    'target_test_case': init_test_case,
                    'test_case_name': init_test_case_name,
                    'original_path': initial_fm_info.test_case_path,
                    'generated_path': generated_path,
                    'Compile_result': 0 if test_status == 'fail_compile' else 1,
                    'Test_result': 0 if test_status != 'success' else 1,
                    "running_result": test_status, 
                    'test_log': test_log
                })

        with open(gen_save_path, 'w') as f:
            json.dump(total_generated_test_cases, f, indent=4)
    
    os.system(f'rm -rf {configs.project_path_no_test_file}')


def prepare_save_path(save_path):
    os.makedirs(os.path.dirname(save_path), exist_ok=True)

    gen_save_path = f'{save_path[:-5]}'
    
    suffix = ''
    if args.start_at > 0:
        suffix += f'_start_{args.start_at}'

    if args.specify_ids:
        suffix += f'_specify'

    suffix += '.json'

    gen_save_path = gen_save_path + suffix

    return gen_save_path


def check_environment(test_case_and_paths):
    for tc, tc_path in test_case_and_paths:
        print(f'[INFO] Checking the environment for {tc_path}...')
        os.makedirs(os.path.dirname(tc_path), exist_ok=True)
        with open(tc_path, 'w') as f:
            f.write(tc)

        cwd_path = tc_path.split('/src/test/')[0]
        mvn_compile_cmd = ['mvn', 'clean', 'test-compile']
        compile_result = subprocess.run(mvn_compile_cmd, cwd=cwd_path, stdout=subprocess.PIPE, stderr=subprocess.PIPE, universal_newlines=True)

        os.remove(tc_path)

        if 'BUILD SUCCESS' not in compile_result.stdout:
            print(f'[ERROR] {tc_path} cannot be compiled.')
            print(compile_result.stdout)
            print(compile_result.stderr)
            exit(1)


def clean_test_cases(test_case_paths):
    for test_case_path in test_case_paths:
        if os.path.exists(test_case_path):
            os.system(f'rm {test_case_path}')


def get_coverage_info(configs, target_test_case_rel_path, focal_method_name_parameter, focal_file_path):
    test_runner = TestCaseRunner(configs, configs.test_case_run_log_dir)
    test_case_path = f"{configs.project_dir_no_test_file}/{target_test_case_rel_path}"
    focal_file_coverage, fm_cov_statistic_by_jacoco = test_runner.get_coverage_jacoco(test_case_path, focal_file_path, focal_method_name_parameter)
    return focal_file_coverage, fm_cov_statistic_by_jacoco



def main():
    generation_pipeline()


if __name__ == '__main__':
    parser = argparse.ArgumentParser()
    parser.add_argument('--project_name', type=str)
    parser.add_argument('--llm_name', type=str, default='gpt-o4-mini', choices=['gpt-o4-mini', 'deepseek-v3.1'])
    parser.add_argument('--start_at', type=int, default=0)
    parser.add_argument('--stop_at', type=int, default=999999)
    parser.add_argument('--specify_ids', type=lambda s: [int(x) for x in s.split(',')], default=[])
    parser.add_argument('--junit_version', type=str, default='4')

    args = parser.parse_args()

    configs = Configs(args.project_name, args.llm_name)

    generated_test_case_save_path = f'./data/baseline/vanilla/generated_test_cases/{args.llm_name}/{args.project_name}.json'
    os.makedirs(os.path.dirname(generated_test_case_save_path), exist_ok=True)
    configs.generated_test_case_save_path = generated_test_case_save_path

    print(f'Args:\n{args}\n\n')
    print(f'Configs:\n{configs.__dict__}\n\n')

    print(f"Processing {configs.project_name}...\n\n")

    main()