testcases/test_main_pipeline_eval.py

#!/usr/bin/env python3
"""
Simple test function to run main.py with cadillac_testcase config
and validate that specific geometries have R² scores greater than 0.7.
"""


########################################## global variables #########################################################

config_path = "/raid/ansysai/udbhav/alphaLPFM/testcases/config_eval.yaml"
main_script_path = "/raid/ansysai/udbhav/alphaLPFM/main.py"
project_root = "/raid/ansysai/udbhav/alphaLPFM"

test_geometries = ['cadillac_19', 'cadillac_73', 'cadillac_104']
thresholds = {
    'cadillac_19': 0.895,
    'cadillac_73': 0.855,
    'cadillac_104': 0.769,
}



###### import libraries #######
import pytest
import os
import sys
import subprocess
import tempfile
import shutil
from pathlib import Path
from omegaconf import OmegaConf


class TestMainPipeline:
    """Test class for running main.py and validating R² scores."""
    
    @pytest.fixture(scope="class")
    def config_path(self):
        """Path to the cadillac_testcase config file."""
        # return "/raid/ansysai/udbhav/alphaLPFM/configs/testcase/cadillac_testcase/config.yaml"
        return config_path
    
    @pytest.fixture(scope="class")
    def main_script_path(self):
        """Path to the main.py script."""
        return main_script_path
    
    @pytest.fixture(scope="class")
    def project_root(self):
        """Project root directory."""
        return project_root
    
    def test_main_pipeline_r2_scores(self, config_path, main_script_path, project_root):
        """
        Test that main.py runs successfully and top 5 R² scores are > test_r2_threshold.
        
        This test:
        1. Runs main.py with cadillac_testcase config
        2. Loads the R² scores from the results
        3. Validates that top 5 scores are > test_r2_threshold
        """
        # Verify config file exists
        assert os.path.exists(config_path), f"Config file not found: {config_path}"
        assert os.path.exists(main_script_path), f"Main script not found: {main_script_path}"
        
        # Load config to get test_name for path construction
        cfg = OmegaConf.load(config_path)
        test_name = cfg.test_name
        project_name = cfg.project_name
        model_name = cfg.model
        
        # Construct expected output path
        expected_output_dir = os.path.join(
            project_root, 
            "metrics", 
            project_name, 
            f"{model_name}_{test_name}"
        )
        
        print(f"Running main.py with config: {config_path}")
        print(f"Expected output directory: {expected_output_dir}")
        
        # Change to project root directory
        original_cwd = os.getcwd()
        os.chdir(project_root)
        
        try:
            # Run main.py with the config
            cmd = [sys.executable, "main.py", "--config_path", config_path]
            print(f"Executing command: {' '.join(cmd)}")
            
            # Run the command
            result = subprocess.run(
                cmd, 
                capture_output=True, 
                text=True, 
                timeout=1800  # 30 minutes timeout
            )
            
            # Check if command succeeded
            if result.returncode != 0:
                print(f"STDOUT: {result.stdout}")
                print(f"STDERR: {result.stderr}")
                pytest.fail(f"main.py failed with return code {result.returncode}")
            
            print("✅ main.py executed successfully")
            print(f"STDOUT: {result.stdout}")
            
            # Verify output directory was created
            assert os.path.exists(expected_output_dir), f"Output directory not created: {expected_output_dir}"
            
            # Check for best_case directory
            best_case_dir = os.path.join(expected_output_dir, "best_case")
            assert os.path.exists(best_case_dir), f"Best case directory not found: {best_case_dir}"
            
            # Check for test results directory
            test_dir = os.path.join(best_case_dir, "test")
            assert os.path.exists(test_dir), f"Test directory not found: {test_dir}"
            
            # Check for R² scores file
            r_squared_file = os.path.join(test_dir, "r_squared.txt")
            assert os.path.exists(r_squared_file), f"R² scores file not found: {r_squared_file}"
            
            # Load and validate R² scores
            r2_scores_dict = self._load_r2_scores_with_ids(r_squared_file)
            print(f"Loaded {len(r2_scores_dict)} R² scores")
            
            # Specific geometries to test with required thresholds
            
            print(f"Testing specific geometries: {test_geometries} with thresholds: {thresholds}")
            
            # Validate that specific geometries have R² > 0.7
            failed_geometries = []
            passed_geometries = []
            
            for geometry_id in test_geometries:
                if geometry_id in r2_scores_dict:
                    r2_score = r2_scores_dict[geometry_id]
                    required = thresholds[geometry_id]
                    if r2_score > 0.99*required and r2_score < 1.01*required:
                        passed_geometries.append((geometry_id, r2_score))
                        print(f"✅ {geometry_id}: R² = {r2_score:.6f} > {0.99*required} and {r2_score:.6f} < {1.01*required}")
                    else:
                        failed_geometries.append((geometry_id, r2_score))
                        print(f"❌ {geometry_id}: R² = {r2_score:.6f} <= {0.99*required} or {r2_score:.6f} >= {1.01*required}")
                else:
                    failed_geometries.append((geometry_id, "NOT_FOUND"))
                    print(f"❌ {geometry_id}: Not found in results")
            
            if failed_geometries:
                failed_msg = "Specific geometries validation failed:\n"
                for geometry_id, score in failed_geometries:
                    failed_msg += f"  - {geometry_id}: {score}\n"
                failed_msg += f"\nPassed geometries: {passed_geometries}"
                pytest.fail(failed_msg)
            
            print(f"✅ All specific geometries validation passed!")
            print(f"Passed geometries: {passed_geometries}")
                        

            
        except subprocess.TimeoutExpired:
            pytest.fail("main.py execution timed out after 30 minutes")
        except Exception as e:
            pytest.fail(f"Unexpected error during main.py execution: {str(e)}")
        finally:
            # Restore original working directory
            os.chdir(original_cwd)
    
    def _load_r2_scores_with_ids(self, r_squared_file):
        """
        Load R² scores from the results file with geometry IDs.
        
        Expected format: "geometry_id r2_score"
        Returns: dict with geometry_id as key and r2_score as value
        """
        r2_scores_dict = {}
        
        with open(r_squared_file, 'r') as f:
            for line in f:
                line = line.strip()
                if line and not line.startswith('#'):
                    parts = line.split()
                    if len(parts) >= 2:
                        try:
                            geometry_id = parts[0]
                            r2_score = float(parts[1])
                            r2_scores_dict[geometry_id] = r2_score
                        except ValueError:
                            print(f"Warning: Could not parse R² score from line: {line}")
                            continue
        
        return r2_scores_dict
    

if __name__ == "__main__":
    # Run the tests
    pytest.main([__file__, "-v", "--tb=short"])