Atomizer/optimization_engine/hooks/test_hooks.py

"""
Test script for NX Open hooks.

This script tests the hooks module with a real NX part.
Run with: python -m optimization_engine.hooks.test_hooks
"""

import os
import sys
import json

# Add the project root to path
project_root = os.path.dirname(os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
sys.path.insert(0, project_root)

from optimization_engine.hooks.nx_cad import (
    part_manager,
    expression_manager,
    geometry_query,
    feature_manager,
)


def test_hooks(part_path: str):
    """Test all hooks with the given part."""
    print(f"\n{'='*60}")
    print(f"Testing NX Open Hooks")
    print(f"Part: {part_path}")
    print(f"{'='*60}\n")

    if not os.path.exists(part_path):
        print(f"ERROR: Part file not found: {part_path}")
        return False

    all_passed = True

    # Test 1: Get expressions
    print("\n--- Test 1: Get Expressions ---")
    result = expression_manager.get_expressions(part_path)
    if result["success"]:
        print(f"SUCCESS: Found {result['data']['count']} expressions")
        # Show first 5 expressions
        for i, (name, expr) in enumerate(list(result['data']['expressions'].items())[:5]):
            print(f"  {name} = {expr['value']} {expr.get('units', '')}")
        if result['data']['count'] > 5:
            print(f"  ... and {result['data']['count'] - 5} more")
    else:
        print(f"FAILED: {result['error']}")
        all_passed = False

    # Test 2: Get mass properties
    print("\n--- Test 2: Get Mass Properties ---")
    result = geometry_query.get_mass_properties(part_path)
    if result["success"]:
        data = result['data']
        print(f"SUCCESS:")
        print(f"  Mass: {data['mass']:.6f} {data['mass_unit']}")
        print(f"  Volume: {data['volume']:.2f} {data['volume_unit']}")
        print(f"  Surface Area: {data['surface_area']:.2f} {data['area_unit']}")
        print(f"  Material: {data.get('material', 'N/A')}")
        print(f"  Centroid: ({data['centroid']['x']:.2f}, {data['centroid']['y']:.2f}, {data['centroid']['z']:.2f}) mm")
    else:
        print(f"FAILED: {result['error']}")
        all_passed = False

    # Test 3: Get bodies
    print("\n--- Test 3: Get Bodies ---")
    result = geometry_query.get_bodies(part_path)
    if result["success"]:
        data = result['data']
        print(f"SUCCESS:")
        print(f"  Total bodies: {data['count']}")
        print(f"  Solid bodies: {data['solid_count']}")
        print(f"  Sheet bodies: {data['sheet_count']}")
    else:
        print(f"FAILED: {result['error']}")
        all_passed = False

    # Test 4: Get features
    print("\n--- Test 4: Get Features ---")
    result = feature_manager.get_features(part_path)
    if result["success"]:
        data = result['data']
        print(f"SUCCESS: Found {data['count']} features ({data['suppressed_count']} suppressed)")
        # Show first 5 features
        for i, feat in enumerate(data['features'][:5]):
            status = "suppressed" if feat['is_suppressed'] else "active"
            print(f"  {feat['name']} ({feat['type']}): {status}")
        if data['count'] > 5:
            print(f"  ... and {data['count'] - 5} more")
    else:
        print(f"FAILED: {result['error']}")
        all_passed = False

    # Summary
    print(f"\n{'='*60}")
    if all_passed:
        print("ALL TESTS PASSED!")
    else:
        print("SOME TESTS FAILED")
    print(f"{'='*60}\n")

    return all_passed


def main():
    """Main entry point."""
    # Default to bracket study part
    default_part = os.path.join(
        project_root,
        "studies/bracket_stiffness_optimization_V3/1_setup/model/Bracket.prt"
    )

    # Use command line argument if provided
    if len(sys.argv) > 1:
        part_path = sys.argv[1]
    else:
        part_path = default_part

    success = test_hooks(part_path)
    sys.exit(0 if success else 1)


if __name__ == "__main__":
    main()