mcp_server.py

import os
import json
import subprocess
import yaml
from typing import List, Optional, Dict, Any
from mcp.server.fastmcp import FastMCP

from src.analyzer import CodeAnalyzer
from src.test_parser import TestAnalyzer
from src.strategy import StrategyGenerator
from src.generator import TestGenerator

# Initialize FastMCP server
mcp = FastMCP("CppUnitTestGenerator")

@mcp.tool()
def analyze_source_code(source_file_path: str) -> str:
    """
    Analyzes a C++ source file using libclang to extract functions, types, macros, and dependencies.
    
    Args:
        source_file_path: Absolute path to the C++ source file (.cpp, .c, .hpp, .h).
        
    Returns:
        A JSON string containing the AST analysis results.
    """
    if not os.path.exists(source_file_path):
        return f"Error: File not found at {source_file_path}"
    
    analyzer = CodeAnalyzer()
    try:
        results = analyzer.analyze_file(source_file_path)
        
        # Save to TestStrategy
        strategy_dir = "TestStrategy"
        os.makedirs(strategy_dir, exist_ok=True)
        out_name = os.path.basename(source_file_path) + "_analysis.json"
        out_path = os.path.join(strategy_dir, out_name)
        with open(out_path, "w") as f:
            json.dump(results, f, indent=2)
            
        return json.dumps({
            "analysis_result": results,
            "artifact_saved_at": out_path
        }, indent=2)
    except Exception as e:
        return f"Error analyzing file: {str(e)}"

@mcp.tool()
def analyze_existing_tests(test_file_path: str) -> str:
    """
    Analyzes existing C++ test files (GTest) to extract test cases and their bodies.
    
    Args:
        test_file_path: Absolute path to the test file or directory.
    
    Returns:
        A JSON string containing the list of found tests and their details.
    """
    if not os.path.exists(test_file_path):
        return f"Error: Path not found at {test_file_path}"
    
    test_files = []
    if os.path.isdir(test_file_path):
        for root, _, files in os.walk(test_file_path):
            for f in files:
                if f.endswith((".cpp", ".cc", ".cxx", ".c")):
                    test_files.append(os.path.join(root, f))
    else:
        test_files.append(test_file_path)
    
    test_analyzer = TestAnalyzer()
    all_tests = []
    try:
        for tf in test_files:
            found_tests = test_analyzer.analyze_test_file(tf)
            all_tests.extend(found_tests)
            
        # Save to TestStrategy
        strategy_dir = "TestStrategy"
        os.makedirs(strategy_dir, exist_ok=True)
        
        out_name = "existing_tests_analysis.json"
        if len(test_files) == 1:
            out_name = os.path.basename(test_files[0]) + "_analysis.json"
            
        out_path = os.path.join(strategy_dir, out_name)
        with open(out_path, "w") as f:
            json.dump(all_tests, f, indent=2)
            
        return json.dumps({
            "existing_tests": all_tests,
            "artifact_saved_at": out_path
        }, indent=2)
    except Exception as e:
        return f"Error analyzing tests: {str(e)}"

@mcp.tool()
def get_test_strategy_context(source_file_path: str, test_file_path: Optional[str] = None) -> str:
    """
    Generates a technical test strategy context by comparing source code with existing tests.
    This provides boundary conditions, equivalence partitions, and MCDC requirements.
    
    Args:
        source_file_path: Absolute path to the source file.
        test_file_path: Optional path to existing tests.
    
    Returns:
        A JSON string representing the FileStrategy (technical baseline).
    """
    analyzer = CodeAnalyzer()
    test_analyzer = TestAnalyzer()
    strategy_gen = StrategyGenerator()

    try:
        analysis_results = analyzer.analyze_file(source_file_path)
        
        existing_tests = []
        if test_file_path and os.path.exists(test_file_path):
            if os.path.isdir(test_file_path):
                for root, _, files in os.walk(test_file_path):
                    for f in files:
                        if f.endswith((".cpp", ".cc", ".cxx", ".c")):
                            existing_tests.extend(test_analyzer.analyze_test_file(os.path.join(root, f)))
            else:
                existing_tests.extend(test_analyzer.analyze_test_file(test_file_path))
        
        strategy = strategy_gen.generate_strategy(analysis_results, existing_tests)
        
        # Save artifacts to TestStrategy folder (matching main.py behavior)
        strategy_dir = "TestStrategy"
        if not os.path.exists(strategy_dir):
            os.makedirs(strategy_dir, exist_ok=True)
            
        output_base = os.path.splitext(os.path.basename(source_file_path))[0]
        yaml_path = os.path.join(strategy_dir, f"{output_base}_strategy.yaml")
        md_path = os.path.join(strategy_dir, f"{output_base}_strategy.md")
        
        strategy_gen.save_yaml(strategy, yaml_path)
        strategy_gen.save_markdown(strategy, md_path)

        # Convert dataclass to dict and add guidelines/review instructions
        import dataclasses
        strategy_dict = dataclasses.asdict(strategy)
        return json.dumps({
            **strategy_dict,
            "artifact_paths": {
                "yaml": yaml_path,
                "markdown": md_path
            },
            "user_action_required": "Please REVIEW the test strategy in the TestStrategy folder before proceeding to generation.",
            "optimization_guidelines": [
                "Prioritize Failure Injection: Ensure test scenarios for all identified 'mocks_needed' are created.",
                "Corner Case Focus: Use the provided Boundary Conditions to target edge-of-envelope behavior.",
                "High-Impact Assertions: Focus on verifying design intent and state changes rather than trivial getters.",
                "Mock Verification: Use EXPECT_CALL to verify side-effects and ON_CALL to define failure behaviors."
            ]
        }, indent=2)
    except Exception as e:
        return f"Error generating strategy: {str(e)}"

@mcp.tool()
def generate_gtest_file(strategy_json: str, output_file_path: str) -> str:
    """
    Generates a complete GTest C++ file from a strategy JSON. 
    The strategy JSON should contain 'test_body' for any new tests to be generated.
    
    Args:
        strategy_json: The strategy data as a JSON string.
        output_file_path: Absolute path where the generated test file should be saved.
        
    Returns:
        A success message or error details.
    """
    try:
        strategy_dict = json.loads(strategy_json)
        from src.strategy import FileStrategy, StrategyGenerator
        strategy = FileStrategy.from_dict(strategy_dict)
        
        # Save updated artifacts to TestStrategy folder
        strategy_dir = "TestStrategy"
        os.makedirs(strategy_dir, exist_ok=True)
        
        source_file_path = strategy.source_file
        output_base = os.path.splitext(os.path.basename(source_file_path))[0]
        yaml_path = os.path.join(strategy_dir, f"{output_base}_strategy.yaml")
        md_path = os.path.join(strategy_dir, f"{output_base}_strategy.md")
        
        strategy_gen = StrategyGenerator()
        strategy_gen.save_yaml(strategy, yaml_path)
        strategy_gen.save_markdown(strategy, md_path)
        
        # Now generate the test code using the updated YAML on disk
        generator = TestGenerator()
        if not output_file_path.startswith("GeneratedUT"):
             if not os.path.exists("GeneratedUT"):
                 os.makedirs("GeneratedUT", exist_ok=True)
             output_file_path = os.path.join("GeneratedUT", os.path.basename(output_file_path))
             
        generator.generate_test_code(yaml_path, output_file_path, use_llm=False)
        return f"Successfully generated test file at {output_file_path} and updated strategy in {strategy_dir}. PLEASE OPEN AND REVIEW the generated code before proceeding to build/verification."
    except Exception as e:
        return f"Error generating test code: {str(e)}"

@mcp.tool()
def run_and_verify_tests(source_file_path: str, test_file_path: str, run: bool = True, coverage: bool = True) -> str:
    """
    Generates CMake, builds the project, runs the tests, and reports coverage.
    
    Args:
        source_file_path: Absolute path to the source file.
        test_file_path: Absolute path to the test file.
        run: Whether to run the tests after building.
        coverage: Whether to generate coverage reports.
        
    Returns:
        A text summary of the build, run, and coverage results.
    """
    try:
        generator = TestGenerator()
        # Ensure we are in the project root or a specific build dir
        # For simplicity, we'll use the dir of the source file or current wd
        cwd = os.getcwd()
        
        # 1. Generate CMake in GeneratedUT
        cmake_path = os.path.join("GeneratedUT", "CMakeLists.txt")
        generator.generate_cmake(source_file_path, test_file_path, output_file=cmake_path)
        
        # 2. Build in GeneratedUT/build
        build_dir = os.path.join(cwd, "GeneratedUT", "build")
        if not os.path.exists(build_dir):
            os.makedirs(build_dir)
            
        subprocess.run(["cmake", ".."], cwd=build_dir, check=True, capture_output=True, text=True)
        subprocess.run(["make"], cwd=build_dir, check=True, capture_output=True, text=True)
        
        output = ["Build successful."]
        
        if run:
            proj_name = os.path.splitext(os.path.basename(source_file_path))[0]
            exe_path = f"./{proj_name}_test"
            run_res = subprocess.run([exe_path], cwd=build_dir, capture_output=True, text=True)
            output.append(f"Test Execution Output:\n{run_res.stdout}\n{run_res.stderr}")
            
            if coverage:
                # GCOV
                subprocess.run(["gcov", "-o", f"CMakeFiles/{proj_name}_test.dir/", source_file_path], cwd=build_dir, capture_output=True, text=True)
                
                # LCOV and GenHTML
                coverage_file = os.path.join(cwd, "GeneratedUT", "coverage.info")
                coverage_html_dir = os.path.join(cwd, "GeneratedUT", "coverage_html")
                
                subprocess.run(["lcov", "--capture", "--directory", build_dir, "--output-file", coverage_file], capture_output=True, text=True)
                subprocess.run(["lcov", "--extract", coverage_file, f"*{source_file_path}*", "--output-file", coverage_file], capture_output=True, text=True)
                subprocess.run(["genhtml", coverage_file, "--output-directory", coverage_html_dir], capture_output=True, text=True)
                
                output.append(f"LCOV HTML report generated in GeneratedUT/coverage_html/")

                gcov_file = os.path.basename(source_file_path) + ".gcov"
                gcov_full_path = os.path.join(build_dir, gcov_file)
                if os.path.exists(gcov_full_path):
                    hits = 0
                    total = 0
                    with open(gcov_full_path, "r") as f:
                        for line in f:
                            if ":" in line:
                                count = line.split(":")[0].strip()
                                if count.isdigit():
                                    hits += 1
                                    total += 1
                                elif count == "#####":
                                    total += 1
                    if total > 0:
                        percent = (hits / total) * 100
                        output.append(f"Coverage: {percent:.2f}% ({hits}/{total} lines)")
                    else:
                        output.append("Coverage: 0.00% (No lines found)")
                else:
                    output.append("Warning: .gcov file not generated.")
        
        final_output = "\n".join(output)
        
        # Save report
        strategy_dir = "TestStrategy"
        os.makedirs(strategy_dir, exist_ok=True)
        report_name = os.path.basename(source_file_path) + "_run_report.txt"
        report_path = os.path.join(strategy_dir, report_name)
        with open(report_path, "w") as f:
            f.write(final_output)
            
        final_output += f"\n\nReport saved to: {report_path}"
        
        return final_output
    except subprocess.CalledProcessError as e:
        return f"Process failed with return code {e.returncode}\nStdout: {e.stdout}\nStderr: {e.stderr}"
    except Exception as e:
        return f"Error during verification: {str(e)}"

if __name__ == "__main__":
    mcp.run()