test/test_ttrlvr_azr_integration.py

#!/usr/bin/env python3
"""
TTRLVR + AZR 통합 시스템 테스트 스크립트

주요 컴포넌트들의 단위 테스트 및 통합 테스트를 수행합니다:
1. Task Generator 테스트 (AZR 메타데이터 포함)
2. Complete Pipeline 테스트 (basic_accuracy 업데이트)  
3. Data Converter 테스트 (parquet 저장)
4. Iterative Trainer 테스트 (라운드 관리)
5. 전체 통합 테스트
"""

import os
import sys
import json
import tempfile
import shutil
import unittest
from unittest.mock import Mock, patch, MagicMock
from pathlib import Path

# 경로 설정
sys.path.append('/home/ubuntu/RLVR/TestTime-RLVR-v2')

# TTRLVR 모듈 임포트
from absolute_zero_reasoner.testtime.config import TestTimeConfig, BenchmarkConfig
from absolute_zero_reasoner.testtime.logger import TestTimeLogger
from absolute_zero_reasoner.testtime.task_generator import TestTimeTaskGenerator
from absolute_zero_reasoner.testtime.complete_pipeline import CompleteTestTimePipeline
from test.utils.iterative_trainer import IterativeTrainer


class TestTTRLVRAZRIntegration(unittest.TestCase):
    """TTRLVR + AZR 통합 시스템 테스트"""
    
    def setUp(self):
        """테스트 설정"""
        self.config = TestTimeConfig()
        self.config.model_name = "Qwen/Qwen2.5-7B"
        self.config.max_new_tokens = 256
        self.config.temperature = 0.05
        
        self.logger = TestTimeLogger()
        self.test_dir = tempfile.mkdtemp()
        
        # 테스트용 IPO 트리플 데이터
        self.test_ipo_triples = [
            {
                'id': 'HumanEval_0_triple_0',
                'input': '[1, 2, 3]',
                'actual_output': '[2, 4, 6]',
                'program': 'def test_func(lst):\n    return [x * 2 for x in lst]',
                'full_input_str': 'test_func([1, 2, 3])',
                'source_program_id': 'program_0',
                'ipo_index': 0
            },
            {
                'id': 'HumanEval_0_triple_1', 
                'input': '[4, 5]',
                'actual_output': '[8, 10]',
                'program': 'def test_func(lst):\n    return [x * 2 for x in lst]',
                'full_input_str': 'test_func([4, 5])',
                'source_program_id': 'program_0',
                'ipo_index': 1
            }
        ]
        
    def tearDown(self):
        """테스트 정리"""
        if os.path.exists(self.test_dir):
            shutil.rmtree(self.test_dir)
    
    def test_task_generator_azr_metadata(self):
        """Task Generator의 AZR 메타데이터 생성 테스트"""
        
        task_generator = TestTimeTaskGenerator(self.config, self.logger)
        
        # Task 생성 (round_num 포함)
        problem_id = "HumanEval_0"
        round_num = 3
        tasks = task_generator.generate_tasks(self.test_ipo_triples, problem_id, round_num)
        
        # 기본 구조 검증
        self.assertIn('induction', tasks)
        self.assertIn('deduction', tasks) 
        self.assertIn('abduction', tasks)
        
        # 각 task 타입별 검증
        for task_type, task_list in tasks.items():
            self.assertGreater(len(task_list), 0, f"{task_type} tasks should be generated")
            
            for task in task_list:
                # AZR 메타데이터 검증
                self.assertIn('uid', task)
                self.assertIn('ipo_group_id', task)
                self.assertIn('original_problem_id', task)
                self.assertIn('round', task)
                self.assertIn('extra_info', task)
                self.assertIn('basic_accuracy', task)
                self.assertIn('ground_truth', task)
                
                # 값 검증
                self.assertEqual(task['original_problem_id'], problem_id)
                self.assertEqual(task['round'], round_num)
                self.assertEqual(task['basic_accuracy'], 0.0)  # 초기값
                self.assertIn(problem_id, task['uid'])
                self.assertIn(str(round_num), task['uid'])
                self.assertIn(task_type, task['uid'])
                
                # task 타입별 metric 검증
                if task_type == 'induction':
                    self.assertEqual(task['extra_info']['metric'], 'code_f')
                elif task_type == 'deduction':
                    self.assertEqual(task['extra_info']['metric'], 'code_o')
                elif task_type == 'abduction':
                    self.assertEqual(task['extra_info']['metric'], 'code_i')
        
        print("✅ Task Generator AZR metadata test passed")
    
    def test_data_converter_parquet_format(self):
        """데이터 변환기의 parquet 형식 테스트"""
        
        # Mock task 데이터 생성
        mock_tasks = {
            'induction': [
                {
                    'task_id': 'induction_0',
                    'task_type': 'induction', 
                    'prompt': 'Test prompt',
                    'uid': 'HumanEval_0_round_1_induction_0',
                    'ipo_group_id': 'HumanEval_0_program_0_ipo_0',
                    'source_program_id': 'program_0',
                    'ipo_index': 0,
                    'ipo_triple': {
                        'input': '[1, 2, 3]',
                        'output': '[2, 4, 6]',
                        'program': 'def test_func(lst):\n    return [x * 2 for x in lst]'
                    },
                    'ground_truth': 'def test_func(lst):\n    return [x * 2 for x in lst]',
                    'extra_info': {'metric': 'code_f'},
                    'basic_accuracy': 1.0,
                    'original_problem_id': 'HumanEval_0',
                    'round': 1
                }
            ]
        }
        
        # Complete pipeline mock 생성
        with patch('absolute_zero_reasoner.testtime.complete_pipeline.CompleteTestTimePipeline') as mock_pipeline:
            pipeline = CompleteTestTimePipeline(self.config, self.logger)
            
            # _save_azr_training_data 메서드 테스트
            output_dir = self.test_dir
            problem_id = "HumanEval_0"
            round_num = 1
            
            saved_files = pipeline._save_azr_training_data(mock_tasks, problem_id, round_num, output_dir)
            
            # 파일 생성 검증
            self.assertIn('induction', saved_files)
            self.assertTrue(os.path.exists(saved_files['induction']))
            
            # Parquet 파일 읽기 테스트
            import pandas as pd
            df = pd.read_parquet(saved_files['induction'])
            
            # 데이터 검증
            self.assertEqual(len(df), 1)
            self.assertIn('prompt', df.columns)
            self.assertIn('uid', df.columns)
            self.assertIn('ipo_group_id', df.columns)
            self.assertIn('ground_truth', df.columns)
            self.assertIn('basic_accuracy', df.columns)
            
            # 프롬프트 형식 검증 (chat 형식)
            prompt_data = df.iloc[0]['prompt']
            self.assertIsInstance(prompt_data, list)
            self.assertEqual(prompt_data[0]['role'], 'user')
            self.assertIn('content', prompt_data[0])
        
        print("✅ Data converter parquet format test passed")
    
    def test_complete_pipeline_basic_accuracy_update(self):
        """Complete Pipeline의 basic_accuracy 업데이트 테스트"""
        
        # Mock components
        with patch.multiple(
            'absolute_zero_reasoner.testtime.complete_pipeline.CompleteTestTimePipeline',
            _generate_task_response=Mock(return_value="test response"),
            _extract_answer_by_task_type=Mock(return_value="test answer"),
            _calculate_task_accuracy=Mock(return_value=0.8)
        ):
            pipeline = CompleteTestTimePipeline(self.config, self.logger)
            
            # Mock task 데이터
            mock_tasks = {
                'induction': [
                    {
                        'task_id': 'induction_0',
                        'prompt': 'test prompt',
                        'expected_solution': 'test solution',
                        'evaluation_data': {'test': 'data'},
                        'basic_accuracy': 0.0  # 초기값
                    }
                ]
            }
            
            # Task 평가 실행
            evaluations = pipeline._evaluate_tasks_with_llm(mock_tasks)
            
            # basic_accuracy 업데이트 검증
            updated_task = mock_tasks['induction'][0]
            self.assertEqual(updated_task['basic_accuracy'], 0.8)  # Mock에서 반환한 값
            
            # Evaluation 결과 검증
            self.assertIn('induction', evaluations)
            eval_result = evaluations['induction'][0]
            self.assertIn('basic_accuracy', eval_result)
            self.assertEqual(eval_result['basic_accuracy'], 0.8)
        
        print("✅ Complete pipeline basic_accuracy update test passed")
    
    def test_iterative_trainer_round_management(self):
        """Iterative Trainer의 라운드 관리 테스트"""
        
        # Mock benchmark config
        benchmark_config = BenchmarkConfig(
            name='test_benchmark',
            problems_path='/test/path',
            max_problems=None
        )
        
        problem_ids = ['TestProblem_1', 'TestProblem_2']
        
        with patch.object(IterativeTrainer, '_update_pipeline_model'):
            with patch.object(IterativeTrainer, '_train_azr_with_round_data') as mock_train:
                with patch.object(IterativeTrainer, '_save_checkpoint'):
                    with patch('absolute_zero_reasoner.testtime.complete_pipeline.CompleteTestTimePipeline') as mock_pipeline_class:
                        
                        # Mock pipeline 결과
                        mock_pipeline = Mock()
                        mock_pipeline.run_complete_pipeline.return_value = {
                            'success': True,
                            'azr_training_data': {'induction': '/test/path/induction.parquet'},
                            'steps': {
                                'azr_data_saving': {'total_tasks': 10}
                            }
                        }
                        mock_pipeline_class.return_value = mock_pipeline
                        
                        # Mock AZR 학습 결과
                        mock_train.return_value = "/data/RLVR/checkpoints/ttrlvr_azr/round_1"
                        
                        # Trainer 초기화
                        trainer = IterativeTrainer(self.config, self.logger)
                        trainer.checkpoint_dir = self.test_dir
                        
                        # 단일 라운드 테스트 (전체 30라운드는 너무 오래 걸림)
                        round_result = trainer._run_single_round(benchmark_config, problem_ids, 1)
                        
                        # 결과 검증
                        self.assertTrue(round_result['success'])
                        self.assertEqual(len(round_result['problems']), 2)
                        self.assertGreater(len(round_result['training_data_files']), 0)
                        
                        # 통계 검증
                        stats = round_result['stats']
                        self.assertEqual(stats['total_problems'], 2)
                        self.assertEqual(stats['successful_problems'], 2)
                        self.assertEqual(stats['failed_problems'], 0)
        
        print("✅ Iterative trainer round management test passed")
    
    def test_data_combination_and_sorting(self):
        """라운드 데이터 통합 및 정렬 테스트"""
        
        trainer = IterativeTrainer(self.config, self.logger)
        
        # Mock training data files
        training_data_files = [
            {
                'problem_id': 'TestProblem_1',
                'files': {
                    'induction': os.path.join(self.test_dir, 'test_induction_1.parquet')
                }
            },
            {
                'problem_id': 'TestProblem_2', 
                'files': {
                    'induction': os.path.join(self.test_dir, 'test_induction_2.parquet')
                }
            }
        ]
        
        # Mock parquet 파일 생성
        import pandas as pd
        
        # 첫 번째 문제 데이터
        data1 = pd.DataFrame([
            {
                'uid': 'TestProblem_1_round_1_induction_0',
                'ipo_group_id': 'TestProblem_1_program_1_ipo_2',
                'basic_accuracy': 0.8
            }
        ])
        data1.to_parquet(training_data_files[0]['files']['induction'], index=False)
        
        # 두 번째 문제 데이터
        data2 = pd.DataFrame([
            {
                'uid': 'TestProblem_2_round_1_induction_0',
                'ipo_group_id': 'TestProblem_2_program_0_ipo_1', 
                'basic_accuracy': 0.6
            }
        ])
        data2.to_parquet(training_data_files[1]['files']['induction'], index=False)
        
        # 데이터 통합 테스트
        combined_path = trainer._combine_round_data(training_data_files, 1)
        
        self.assertIsNotNone(combined_path)
        self.assertTrue(os.path.exists(combined_path))
        
        # 통합된 파일 검증
        combined_file = os.path.join(combined_path, 'induction.parquet')
        self.assertTrue(os.path.exists(combined_file))
        
        # 데이터 정렬 검증 (ipo_group_id로 정렬되어야 함)
        combined_df = pd.read_parquet(combined_file)
        self.assertEqual(len(combined_df), 2)
        
        # ipo_group_id 정렬 검증
        ipo_groups = combined_df['ipo_group_id'].tolist()
        self.assertEqual(ipo_groups, sorted(ipo_groups))
        
        print("✅ Data combination and sorting test passed")


class TestPerformanceAndMemory(unittest.TestCase):
    """성능 및 메모리 테스트"""
    
    def setUp(self):
        self.config = TestTimeConfig()
        self.logger = TestTimeLogger()
    
    def test_memory_cleanup_between_rounds(self):
        """라운드 간 메모리 정리 테스트"""
        
        with patch('absolute_zero_reasoner.testtime.complete_pipeline.CompleteTestTimePipeline') as mock_pipeline_class:
            mock_pipeline = Mock()
            mock_pipeline.model = Mock()
            mock_pipeline.tokenizer = Mock()
            mock_pipeline_class.return_value = mock_pipeline
            
            trainer = IterativeTrainer(self.config, self.logger)
            
            # 모델 업데이트 테스트 (메모리 정리 포함)
            trainer._update_pipeline_model("/new/model/path")
            
            # 모델과 토크나이저가 None으로 설정되었는지 확인
            self.assertIsNone(trainer.complete_pipeline.model)
            self.assertIsNone(trainer.complete_pipeline.tokenizer)
        
        print("✅ Memory cleanup between rounds test passed")
    
    def test_checkpoint_size_and_structure(self):
        """체크포인트 크기 및 구조 테스트"""
        
        trainer = IterativeTrainer(self.config, self.logger)
        test_dir = tempfile.mkdtemp()
        trainer.checkpoint_dir = test_dir
        
        try:
            # Mock training results
            training_results = {
                'total_rounds': 30,
                'rounds': {
                    1: {'success': True, 'stats': {'total_tasks': 100}},
                    2: {'success': True, 'stats': {'total_tasks': 95}},
                    3: {'success': False, 'error': 'Test error'}
                }
            }
            
            # 체크포인트 저장
            trainer._save_checkpoint(3, "/test/model/path", training_results)
            
            # 체크포인트 파일 존재 확인
            checkpoint_path = os.path.join(test_dir, "checkpoint_round_3")
            self.assertTrue(os.path.exists(checkpoint_path))
            
            checkpoint_file = os.path.join(checkpoint_path, "checkpoint.json")
            summary_file = os.path.join(checkpoint_path, "summary.txt")
            
            self.assertTrue(os.path.exists(checkpoint_file))
            self.assertTrue(os.path.exists(summary_file))
            
            # 체크포인트 내용 검증
            with open(checkpoint_file, 'r') as f:
                checkpoint_data = json.load(f)
            
            self.assertEqual(checkpoint_data['round_num'], 3)
            self.assertEqual(checkpoint_data['model_path'], "/test/model/path")
            self.assertIn('training_results', checkpoint_data)
            
            # 파일 크기 확인 (너무 크지 않은지)
            checkpoint_size = os.path.getsize(checkpoint_file)
            self.assertLess(checkpoint_size, 1024 * 1024)  # 1MB 미만
            
        finally:
            shutil.rmtree(test_dir)
        
        print("✅ Checkpoint size and structure test passed")


def run_integration_test():
    """통합 테스트 실행"""
    
    print("🧪 TTRLVR + AZR 통합 시스템 테스트 시작")
    print("=" * 60)
    
    # 테스트 스위트 생성
    loader = unittest.TestLoader()
    suite = unittest.TestSuite()
    
    # 테스트 클래스 추가
    suite.addTests(loader.loadTestsFromTestCase(TestTTRLVRAZRIntegration))
    suite.addTests(loader.loadTestsFromTestCase(TestPerformanceAndMemory))
    
    # 테스트 실행
    runner = unittest.TextTestRunner(verbosity=2)
    result = runner.run(suite)
    
    # 결과 요약
    print("\n" + "=" * 60)
    print("📊 테스트 결과 요약:")
    print(f"   - 실행된 테스트: {result.testsRun}")
    print(f"   - 성공: {result.testsRun - len(result.failures) - len(result.errors)}")
    print(f"   - 실패: {len(result.failures)}")
    print(f"   - 오류: {len(result.errors)}")
    
    if result.failures:
        print("\n❌ 실패한 테스트:")
        for test, traceback in result.failures:
            print(f"   - {test}: {traceback}")
    
    if result.errors:
        print("\n💥 오류 발생 테스트:")
        for test, traceback in result.errors:
            print(f"   - {test}: {traceback}")
    
    success = len(result.failures) == 0 and len(result.errors) == 0
    
    if success:
        print("\n🎉 모든 테스트 성공!")
    else:
        print("\n⚠️ 일부 테스트 실패 또는 오류 발생")
    
    return success


if __name__ == '__main__':
    success = run_integration_test()
    sys.exit(0 if success else 1)