Atomizer/tests/test_code_generation.py

"""
Test Feature Code Generation Pipeline

This test demonstrates the Research Agent's ability to:
1. Learn from a user-provided example (XML material file)
2. Extract schema and patterns
3. Design a feature specification
4. Generate working Python code from the learned template
5. Save the generated code to a file

Author: Atomizer Development Team
Version: 0.1.0 (Phase 2 Week 2)
Last Updated: 2025-01-16
"""

import sys
from pathlib import Path

# Set UTF-8 encoding for Windows console
if sys.platform == 'win32':
    import codecs
    sys.stdout = codecs.getwriter('utf-8')(sys.stdout.buffer, errors='replace')
    sys.stderr = codecs.getwriter('utf-8')(sys.stderr.buffer, errors='replace')

# Add project root to path
project_root = Path(__file__).parent.parent
sys.path.insert(0, str(project_root))

from optimization_engine.research_agent import (
    ResearchAgent,
    ResearchFindings,
    CONFIDENCE_LEVELS
)


def test_code_generation():
    """Test complete code generation workflow from example to working code."""
    print("\n" + "="*80)
    print("FEATURE CODE GENERATION TEST")
    print("="*80)

    agent = ResearchAgent()

    # Step 1: User provides material XML example
    print("\n" + "-"*80)
    print("[Step 1] User Provides Example Material XML")
    print("-"*80)

    example_xml = """<?xml version="1.0" encoding="UTF-8"?>
<PhysicalMaterial name="Steel_AISI_1020" version="1.0">
    <Density units="kg/m3">7850</Density>
    <YoungModulus units="GPa">200</YoungModulus>
    <PoissonRatio>0.29</PoissonRatio>
    <ThermalExpansion units="1/K">1.17e-05</ThermalExpansion>
    <YieldStrength units="MPa">295</YieldStrength>
</PhysicalMaterial>"""

    print("\n  Example XML (steel material):")
    for line in example_xml.split('\n')[:4]:
        print(f"    {line}")
    print("    ...")

    # Step 2: Agent learns from example
    print("\n" + "-"*80)
    print("[Step 2] Agent Learns Schema from Example")
    print("-"*80)

    findings = ResearchFindings(
        sources={'user_example': 'steel_material.xml'},
        raw_data={'user_example': example_xml},
        confidence_scores={'user_example': CONFIDENCE_LEVELS['user_validated']}
    )

    knowledge = agent.synthesize_knowledge(findings)

    print(f"\n  Learned schema:")
    if knowledge.schema and 'xml_structure' in knowledge.schema:
        xml_schema = knowledge.schema['xml_structure']
        print(f"    Root element: {xml_schema['root_element']}")
        print(f"    Attributes: {xml_schema.get('attributes', {})}")
        print(f"    Required fields ({len(xml_schema['required_fields'])}):")
        for field in xml_schema['required_fields']:
            print(f"      - {field}")
    print(f"\n  Confidence: {knowledge.confidence:.2f}")

    # Step 3: Design feature specification
    print("\n" + "-"*80)
    print("[Step 3] Design Feature Specification")
    print("-"*80)

    feature_name = "nx_material_generator"
    feature_spec = agent.design_feature(knowledge, feature_name)

    print(f"\n  Feature designed:")
    print(f"    Feature ID: {feature_spec['feature_id']}")
    print(f"    Category: {feature_spec['category']}")
    print(f"    Subcategory: {feature_spec['subcategory']}")
    print(f"    Lifecycle stage: {feature_spec['lifecycle_stage']}")
    print(f"    Implementation file: {feature_spec['implementation']['file_path']}")
    print(f"    Number of inputs: {len(feature_spec['interface']['inputs'])}")
    print(f"\n  Input parameters:")
    for input_param in feature_spec['interface']['inputs']:
        print(f"    - {input_param['name']}: {input_param['type']}")

    # Step 4: Generate Python code
    print("\n" + "-"*80)
    print("[Step 4] Generate Python Code from Learned Template")
    print("-"*80)

    generated_code = agent.generate_feature_code(feature_spec, knowledge)

    print(f"\n  Generated {len(generated_code)} characters of Python code")
    print(f"\n  Code preview (first 20 lines):")
    print("  " + "-"*76)
    for i, line in enumerate(generated_code.split('\n')[:20]):
        print(f"  {line}")
    print("  " + "-"*76)
    print(f"  ... ({len(generated_code.split(chr(10)))} total lines)")

    # Step 5: Validate generated code
    print("\n" + "-"*80)
    print("[Step 5] Validate Generated Code")
    print("-"*80)

    # Check that code has necessary components
    validations = [
        ('Function definition', f'def {feature_name}(' in generated_code),
        ('Docstring', '"""' in generated_code),
        ('Type hints', ('-> Dict[str, Any]' in generated_code or ': float' in generated_code)),
        ('XML Element handling', 'ET.Element' in generated_code),
        ('Return statement', 'return {' in generated_code),
        ('Example usage', 'if __name__ == "__main__":' in generated_code)
    ]

    all_valid = True
    print("\n  Code validation:")
    for check_name, passed in validations:
        status = "✓" if passed else "✗"
        print(f"    {status} {check_name}")
        if not passed:
            all_valid = False

    assert all_valid, "Generated code is missing required components"

    # Step 6: Save generated code to file
    print("\n" + "-"*80)
    print("[Step 6] Save Generated Code")
    print("-"*80)

    # Create custom_functions directory if it doesn't exist
    custom_functions_dir = project_root / "optimization_engine" / "custom_functions"
    custom_functions_dir.mkdir(parents=True, exist_ok=True)

    output_file = custom_functions_dir / f"{feature_name}.py"
    output_file.write_text(generated_code, encoding='utf-8')

    print(f"\n  Code saved to: {output_file}")
    print(f"  File size: {output_file.stat().st_size} bytes")
    print(f"  Lines of code: {len(generated_code.split(chr(10)))}")

    # Step 7: Test that code is syntactically valid Python
    print("\n" + "-"*80)
    print("[Step 7] Verify Code is Valid Python")
    print("-"*80)

    try:
        compile(generated_code, '<generated>', 'exec')
        print("\n  ✓ Code compiles successfully!")
        print("    Generated code is syntactically valid Python")
    except SyntaxError as e:
        print(f"\n  ✗ Syntax error: {e}")
        assert False, "Generated code has syntax errors"

    # Summary
    print("\n" + "="*80)
    print("CODE GENERATION TEST SUMMARY")
    print("="*80)

    print("\n  Workflow Completed:")
    print("    ✓ User provided example XML")
    print("    ✓ Agent learned schema (5 fields)")
    print("    ✓ Feature specification designed")
    print(f"    ✓ Python code generated ({len(generated_code)} chars)")
    print(f"    ✓ Code saved to {output_file.name}")
    print("    ✓ Code is syntactically valid Python")

    print("\n  What This Demonstrates:")
    print("    - Agent can learn from a single example")
    print("    - Schema extraction works correctly")
    print("    - Code generation follows learned patterns")
    print("    - Generated code has proper structure (docstrings, type hints, examples)")
    print("    - Output is ready to use (valid Python)")

    print("\n  Next Steps (in real usage):")
    print("    1. User tests the generated function")
    print("    2. User provides feedback if adjustments needed")
    print("    3. Agent refines code based on feedback")
    print("    4. Feature gets added to feature registry")
    print("    5. Future requests use this template automatically")

    print("\n" + "="*80)
    print("Code Generation: SUCCESS! ✓")
    print("="*80 + "\n")

    return True


if __name__ == '__main__':
    try:
        success = test_code_generation()
        sys.exit(0 if success else 1)
    except Exception as e:
        print(f"\n[ERROR] {e}")
        import traceback
        traceback.print_exc()
        sys.exit(1)