#!/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)