examiner.py

import os
import re
import torch
from agents import TestExaminationAgent
from configs import Configs
from test_case_runner import TestCaseRunner
from graph_explorer import GraphExplorer
from LSPs.java_lsp import JavaLanguageServer
from parser.java_code_parser import JavaCodeParser
from transformers import AutoModel, AutoTokenizer


class TestExaminer:
    def __init__(self, configs: Configs, max_exploration_depth: int, lsp_server: JavaLanguageServer, skip_deepseek_think: bool = False):
        self.configs = configs
        self.examiner_agent = TestExaminationAgent(
            configs.llm_name, configs.project_name, configs.project_url, 
            configs.openai_key, configs.openai_base_url,
            n_responses=1, skip_deepseek_think=skip_deepseek_think
            )
        self.graph_explorer = GraphExplorer(lsp_server, max_depth=max_exploration_depth, efficieny_mode=True)
        self.java_code_parser = JavaCodeParser()
        self.test_case_runner = TestCaseRunner(configs, configs.test_case_run_log_dir)

        self.embedding_model = AutoModel.from_pretrained("Salesforce/codet5p-110m-embedding", trust_remote_code=True).eval().to('cuda')
        self.tokenizer = AutoTokenizer.from_pretrained("Salesforce/codet5p-110m-embedding", trust_remote_code=True)

    def generate_examination(self, initial_focal_method: str, initial_test: str, focal_method_path: str, test_case_rel_path: str):
        """
        Given an initial focal method and test case, identify the assertions in the initial test case, then generate wrong assertions for each identified assertion. Each identified assertion and its corresponding wrong assertions will be used to generate an examination test case.
        """
        exam_list = []
        # Generate wrong oracles
        oracle_lists = self.examiner_agent.generate_wrong_oracle(initial_focal_method, initial_test)
        if len(oracle_lists) == 0:
            exam_list.append({
                'test_with_exam_oracles': '',
                'exam_oracle_list_str': '',
                'ground_truth': '',
                'all_generated_wrong_oracles': [],
                'valid_wrong_oracles': []
            })
            return exam_list

        for each_oracle_list in oracle_lists:
            original_oracle = each_oracle_list[0]
            wrong_oracles = each_oracle_list[1:]
            # validate each wrong oracle by replacing the original test with each wrong oracle and running the test
            valid_wrong_oracles = []
            for each_wrong_oracle in wrong_oracles:
                test_with_wrong_oracle = self.examiner_agent.replace_oracle(initial_test, original_oracle, each_wrong_oracle)
                is_valid_wrong = self.run_test_case(test_with_wrong_oracle, focal_method_path, test_case_rel_path)
                if is_valid_wrong:
                    valid_wrong_oracles.append(each_wrong_oracle)
                if len(valid_wrong_oracles) == 4:
                    break
            
            if len(valid_wrong_oracles) > 0:
                test_with_exam_oracles, exam_oracle_list_str, ground_truth = self.examiner_agent.generate_examination(initial_test, original_oracle, valid_wrong_oracles)
            else:
                test_with_exam_oracles = ''
                exam_oracle_list_str = ''
                ground_truth = ''
                print(f'[WARNING] No valid wrong oracles found for focal method: {initial_focal_method}, test case: {test_case_rel_path}')

            exam_list.append({
                'test_with_exam_oracles': test_with_exam_oracles,
                'exam_oracle_list_str': exam_oracle_list_str,
                'ground_truth': ground_truth,
                'all_generated_wrong_oracles': wrong_oracles,
                'valid_wrong_oracles': valid_wrong_oracles
            })
        return exam_list
    
    def exmaine_with_fact_collection(self, initial_focal_method: str, focal_file_path: str, focal_method_name: str,
                                     initial_test_case: str, test_case_path: str, test_case_name: str,
                                     exams: list,  max_k_facts_per_round: int = 10, max_round_for_each_exam: int = 5):
        """
        Given an initial focal method and exam list (each exam includes generated test case with examination oracles (i.e., assertions) and the examination oracles list), generate the answer for the examination.
        """
        crucial_facts_scenario = ""
        total_facts_scenario, rest_facts_scenario = [], []
        answer_record = []
        for each_exam_for_oracle in exams:
            test_with_exam_oracles = each_exam_for_oracle['test_with_exam_oracles']
            exam_oracle_list_str = each_exam_for_oracle['exam_oracle_list_str']
            ground_truth = int(each_exam_for_oracle['ground_truth'])
            
            crucial_facts_exam = crucial_facts_scenario
            rest_facts_exam = rest_facts_scenario.copy()
            is_correct = False
            for exam_round in range(max_round_for_each_exam):
                answer = self._generate_answer_with_facts(
                    initial_focal_method, test_with_exam_oracles, exam_oracle_list_str, facts=crucial_facts_exam
                    )

                if type(answer) is int:
                    if answer == ground_truth:
                        is_correct = True
                        break
                    else:
                        print(f'[INFO] Wrong answer: {answer}, Ground truth: {ground_truth}')
                        if exam_round < 2:
                            continue
                        else:
                            answer_question = self.examiner_agent.ask_for_question(
                                initial_focal_method, test_with_exam_oracles, exam_oracle_list_str
                                )
                            answer = self._parse_info_lack(answer_question)
                                
                # collect facts
                if not total_facts_scenario:
                    total_facts_scenario = self.collect_facts(
                        initial_focal_method, focal_file_path, focal_method_name, 
                        initial_test_case, test_case_path, test_case_name
                        )
                    rest_facts_scenario = total_facts_scenario.copy()
                    rest_facts_exam = total_facts_scenario.copy()
                
                # rank facts based on the clues
                ranked_facts_with_score = self.rank_facts(rest_facts_exam, answer)
                select_k_facts = min(
                    max_k_facts_per_round, 
                    sum(1 for each in ranked_facts_with_score if each[1] > 0), 
                    max(len(answer), 3)
                    )
                if select_k_facts == 0:
                    print(f"[WARNING] No candidate facts related to clues: {answer}")
                    crucial_facts_exam = ""
                    rest_facts_exam = []
                    break

                if ranked_facts_with_score[select_k_facts - 1][1] == 0:
                    print(f"\n[WARNING] No candidate facts related to clues.\nClue:\n{answer}.\n\nCurrent Crucial Facts:\n{crucial_facts_exam}\n\nCandidate facts:\n{rest_facts_exam}\n\n")
                crucial_fact_list = [each[0] for each in ranked_facts_with_score[:select_k_facts]]
                rest_facts_exam = [each[0] for each in ranked_facts_with_score[select_k_facts:]]

                # format crucial facts
                for each_crucial_fact in crucial_fact_list:
                    class_name, method_name, method_body, _, method_signature = each_crucial_fact
                    crucial_facts_exam += f"```\n{class_name}{{\n{method_signature}{method_body}\n}}\n```\n\n"  # format: Class{MethodSignature{MethodBody}}
                print(f"\n[INFO] Crucial facts collected:\n{crucial_facts_exam}.\n")

            if is_correct:
                crucial_facts_scenario = crucial_facts_exam
                rest_facts_scenario = rest_facts_exam
                answer_record.append(['correct', answer])
            else:
                answer_record.append(['wrong', answer])

        return answer_record, crucial_facts_scenario, rest_facts_scenario, total_facts_scenario

    def _generate_answer_with_facts(self, initial_focal_method: str, test_with_exam_oracles: str, exam_oracle_list_str: str, facts: str):
        answer = self.examiner_agent.examine(
            initial_focal_method, test_with_exam_oracles, exam_oracle_list_str, facts
            )
        
        if 'LACK NECESSARY INFORMATION' in answer:
            return self._parse_info_lack(answer)
        else:
            if 'MY CHOICE' not in answer:
                print(f"[INFO] Invalid answer format: {answer}.\nRegenerate answer.\n")
                return self._generate_answer_with_facts(initial_focal_method, test_with_exam_oracles, exam_oracle_list_str, facts)
                
            choice = answer.split('MY CHOICE')[1].strip()
            choice = re.findall(r'(\d+)', choice)
            if len(choice) != 1:
                print(f"[INFO] Invalid choice format: {choice}.\nRegenerate answer.\n")
                return self._generate_answer_with_facts(initial_focal_method, test_with_exam_oracles, exam_oracle_list_str, facts)
            
            choice = int(choice[0])
            assert 1 <= choice <= len(exam_oracle_list_str.split('\n')), f"Choice out of range: {choice}"
            return choice

    def _parse_info_lack(self, answer: str):
        lack_info_names = []
        lack_info = answer.split('LACK NECESSARY INFORMATION')[1].strip()
        lack_info = lack_info.split('\n')
        for each_line in lack_info:
            if 'METHOD:' in each_line:
                name = each_line.split('METHOD:')[1].strip()
                info_type = 'method'
            elif 'CONSTRUCTOR:' in each_line:
                name = each_line.split('CONSTRUCTOR:')[1].strip()
                info_type = 'constructor'
            elif 'FIELD:' in each_line:
                name = each_line.split('FIELD:')[1].strip()
                info_type = 'field'
            else:
                continue

            extract_name = re.findall(r'`(.+?)`', name)
            if len(extract_name) == 0:
                extract_name = name.strip()
                print(f"[WARNING] Invalid name in lack information: {each_line}\nFinally use: {extract_name}")
            elif len(extract_name) > 1:
                extract_name = ' '.join(extract_name)
                print(f"[WARNING] Invalid name in lack information: {each_line}\nFinally use: {extract_name}")
            else:
                extract_name = extract_name[0]
            lack_info_names.append((info_type, extract_name))
        return lack_info_names

    def collect_facts(self, initial_focal_method: str, focal_file_path: str, focal_method_name: str, 
                      initial_test_case: str, test_case_path: str, test_case_name: str):
        # collect facts by exploring the focal method and test case based on the relationships
        candidate_facts_fm, _ = self.graph_explorer.explore(f'{self.configs.project_dir_no_test_file}/{focal_file_path}', initial_focal_method, focal_method_name)
        proc_candidate_facts_fm = []   # remove the depth information and process the method names
        for each in candidate_facts_fm:
            each_method_name = each[1].strip()
            if '(' in each[1]:
                each_method_name = each_method_name.split('(')[0].split()[-1]
            proc_candidate_facts_fm.append((each[0], each_method_name, each[2], each[3], each[1]))
        candidate_facts_fm = proc_candidate_facts_fm

        test_case_full_path = f'{self.configs.project_dir_no_test_file}/{test_case_path}'
        os.makedirs(os.path.dirname(test_case_full_path), exist_ok=True)
        with open(test_case_full_path, 'w') as f:
            f.write(initial_test_case)
        try:
            test_method = self.extract_test_method(initial_test_case, test_case_name)
            candidate_facts_tc, _ = self.graph_explorer.explore(test_case_full_path, test_method, test_case_name)
            proc_candidate_facts_tc = []  # remove the depth information and process the method names
            for each in candidate_facts_tc:
                each_method_name = each[1].strip()
                if '(' in each[1]:
                    each_method_name = each_method_name.split('(')[0].split()[-1]
                proc_candidate_facts_tc.append((each[0], each_method_name, each[2], each[3], each[1]))
            candidate_facts_tc = proc_candidate_facts_tc
        except Exception as e:
            os.remove(test_case_full_path)
            raise RuntimeError(f"Failed to explore test case {test_case_path}: {e}")
        os.remove(test_case_full_path)

        candidate_facts = list(set(candidate_facts_fm + candidate_facts_tc))
        filtered_candidate_facts = [each for each in candidate_facts if each[2].strip() not in initial_focal_method and each[2].strip() != '']

        # collect facts by extracting methods and fields from the focal file
        self.java_code_parser.parse_java_file(f'{self.configs.project_dir_no_test_file}/{focal_file_path}')
        constructor_name_body_pairs = self.java_code_parser.get_all_constructor_definition()
        method_name_body_pairs = self.java_code_parser.get_all_method_definition()
        field_definition = self.java_code_parser.get_all_field_definition()
        candidate_facts_ff = []
        focal_class_name = focal_file_path.split('/')[-1].replace('.java', '')
        for each in constructor_name_body_pairs:
            candidate_facts_ff.append((
                focal_class_name, 
                each['name'].text.decode('utf8'), 
                each['body'].text.decode('utf8'), 
                f'{self.configs.project_dir_no_test_file}/{focal_file_path}',
                each['name'].parent.text.decode('utf8').split('{')[0].strip()
                ))
            
        checked_fm = False
        for each in method_name_body_pairs:
            if each['name'].text.decode('utf8') != focal_method_name:
                candidate_facts_ff.append((
                    focal_class_name, 
                    each['name'].text.decode('utf8'), 
                    each['body'].text.decode('utf8'), 
                    f'{self.configs.project_dir_no_test_file}/{focal_file_path}',
                    each['name'].parent.text.decode('utf8').split('{')[0].strip()
                    ))
            else:
                checked_fm = True
        assert checked_fm, f"Focal method {focal_method_name} not found in the focal file {focal_file_path}."

        for each in field_definition:
            candidate_facts_ff.append((each[0], each[1], each[2], each[3], each[0]))

        # merge all candidate facts
        facts = filtered_candidate_facts
        existing_fact_class_method_pairs = [(each[0].strip(), each[4].strip()) for each in facts]
        for each_fact in candidate_facts_ff:
            if each_fact[2].strip() == '':
                continue
            
            if (each_fact[0].strip(), each_fact[4].strip()) in existing_fact_class_method_pairs:
                continue

            facts.append(each_fact)

        return facts[::-1]
    
    def extract_test_method(self, test_case: str, test_case_name: str):
        test_case_lines = test_case.split('\n')
        start_idx = None
        for idx, each_line in enumerate(test_case_lines):
            if each_line.strip().startswith('@Test') or each_line.strip().startswith('@ParameterizedTest'):
                for continue_idx in range(idx+1, len(test_case_lines)):  # sometimes, there are other annotations before the test method, such as @GwtIncompatible("Pattern")
                    if test_case_lines[continue_idx].strip().startswith('@'):
                        continue
                    else:
                        start_idx = continue_idx
                        break
                if test_case_name in test_case_lines[start_idx]:
                    break
                else:
                    start_idx = None
                    continue
        
        if start_idx is None:
            if len([each for each in test_case_lines if '@Test' == each.strip() or '@ParameterizedTest' == each.strip()]) == 1:
                possible_line_ids = []
                for line_idx, each_line in enumerate(test_case_lines):
                    if each_line.strip().startswith('public void') and f' {test_case_name}(' in each_line:
                        possible_line_ids.append(line_idx)
                if len(possible_line_ids) == 1:
                    start_idx = possible_line_ids[0]
                        
        assert start_idx is not None, f"Test method for {test_case_name} not found in the test case."
        test_method = '\n'.join(test_case_lines[start_idx:]).strip()
        return test_method

    def rank_facts(self, facts: list, clues: list):
        # rank the candidate facts based on the clues
        clues_with_emb = []
        for each_clue in clues:
            tokenized_input = self.tokenizer(each_clue[1], return_tensors='pt', ).to('cuda')
            clues_with_emb.append((each_clue[0], self.embedding_model(**tokenized_input)[0]))

        facts_with_score = []
        for each_fact in facts:
            class_name, method_name, method_body, _, _ = each_fact
            if method_body == '':
                continue
            
            class_name_tokenized = self.tokenizer(class_name, return_tensors='pt').to('cuda')
            class_name_emb = self.embedding_model(**class_name_tokenized)[0]
            
            method_name_tokenized = self.tokenizer(method_name, return_tensors='pt').to('cuda')
            method_name_emb = self.embedding_model(**method_name_tokenized)[0]
            
            method_body_tokenized = self.tokenizer(method_body, return_tensors='pt', truncation=True).to('cuda')
            method_body_emb = self.embedding_model(**method_body_tokenized)[0]

            score = 0.0
            for each_clue in clues_with_emb:
                if each_clue[0] == 'method':
                    cos_sim = torch.cosine_similarity(method_name_emb, each_clue[1], dim=0)
                elif each_clue[0] == 'constructor':
                    cos_sim = (torch.cosine_similarity(class_name_emb, each_clue[1], dim=0) + torch.cosine_similarity(method_name_emb, each_clue[1], dim=0)) / 2
                elif each_clue[0] == 'field':
                    cos_sim = max(torch.cosine_similarity(method_name_emb, each_clue[1], dim=0), torch.cosine_similarity(method_body_emb, each_clue[1], dim=0))
                
                score += cos_sim.cpu().item()

            facts_with_score.append((each_fact, score))
        facts_with_score.sort(key=lambda x: x[1], reverse=True)

        return  facts_with_score

    def run_test_case(self, test_case: str, focal_method_path: str, test_case_rel_path: str):
        tc_path = f"{self.configs.project_with_test_workspace}/{test_case_rel_path}"
        os.makedirs(os.path.dirname(tc_path), exist_ok=True)
        with open(tc_path, 'w') as f:
            f.write(test_case)

        log_file_path = self.test_case_runner.run_test_case(
            tc_path, focal_method_path, 'no_ref'
        )
        os.remove(tc_path)
        
        with open(log_file_path, 'r') as f:
            log_content = f.read()
        if 'Tests run: 1, Failures: 1, Errors: 0, Skipped: 0' in log_content:
            return True
        else:
            return False