cluster4npu/tests/test_postprocess_mode.py

#!/usr/bin/env python3
"""
Test script for postprocessing mode switching and visualization.
"""

import sys
import os
current_dir = os.path.dirname(os.path.abspath(__file__))
parent_dir = os.path.dirname(current_dir)
sys.path.append(parent_dir)

from core.nodes.exact_nodes import ExactPostprocessNode

def test_postprocess_node():
    """Test the ExactPostprocessNode for mode switching and configuration."""
    
    print("=== Testing ExactPostprocessNode Mode Switching ===")
    
    # Create node instance
    try:
        node = ExactPostprocessNode()
        print("✓ ExactPostprocessNode created successfully")
        
        # Check if NodeGraphQt is available
        if not hasattr(node, 'set_property'):
            print("⚠ NodeGraphQt not available - using mock properties")
            return True  # Skip tests that require NodeGraphQt
            
    except Exception as e:
        print(f"✗ Error creating node: {e}")
        return False
    
    # Test different postprocessing modes
    test_modes = [
        ('fire_detection', 'No Fire,Fire'),
        ('yolo_v3', 'person,car,bicycle,motorbike,aeroplane'),
        ('yolo_v5', 'person,bicycle,car,motorbike,bus,truck'),
        ('classification', 'cat,dog,bird,fish'),
        ('raw_output', '')
    ]
    
    print("\n--- Testing Mode Switching ---")
    for mode, class_names in test_modes:
        try:
            # Set properties for this mode
            node.set_property('postprocess_type', mode)
            node.set_property('class_names', class_names)
            node.set_property('confidence_threshold', 0.6)
            node.set_property('nms_threshold', 0.4)
            
            # Get configuration
            config = node.get_postprocessing_config()
            options = node.get_multidongle_postprocess_options()
            
            print(f"✓ Mode: {mode}")
            print(f"  - Class names: {class_names}")
            print(f"  - Config: {config['postprocess_type']}")
            if options:
                print(f"  - PostProcessor options created successfully")
            else:
                print(f"  - Warning: PostProcessor options not available")
            
        except Exception as e:
            print(f"✗ Error testing mode {mode}: {e}")
            return False
    
    # Test validation
    print("\n--- Testing Configuration Validation ---")
    try:
        # Valid configuration
        node.set_property('postprocess_type', 'fire_detection')
        node.set_property('confidence_threshold', 0.7)
        node.set_property('nms_threshold', 0.3)
        node.set_property('max_detections', 50)
        
        is_valid, message = node.validate_configuration()
        if is_valid:
            print("✓ Valid configuration passed validation")
        else:
            print(f"✗ Valid configuration failed: {message}")
            return False
        
        # Invalid configuration
        node.set_property('confidence_threshold', 1.5)  # Invalid value
        is_valid, message = node.validate_configuration()
        if not is_valid:
            print(f"✓ Invalid configuration caught: {message}")
        else:
            print("✗ Invalid configuration not caught")
            return False
        
    except Exception as e:
        print(f"✗ Error testing validation: {e}")
        return False
    
    # Test display properties
    print("\n--- Testing Display Properties ---")
    try:
        display_props = node.get_display_properties()
        expected_props = ['postprocess_type', 'class_names', 'confidence_threshold']
        
        for prop in expected_props:
            if prop in display_props:
                print(f"✓ Display property found: {prop}")
            else:
                print(f"✗ Missing display property: {prop}")
                return False
                
    except Exception as e:
        print(f"✗ Error testing display properties: {e}")
        return False
    
    # Test business properties
    print("\n--- Testing Business Properties ---")
    try:
        business_props = node.get_business_properties()
        print(f"✓ Business properties retrieved: {len(business_props)} properties")
        
        # Check key properties exist
        key_props = ['postprocess_type', 'class_names', 'confidence_threshold', 'nms_threshold']
        for prop in key_props:
            if prop in business_props:
                print(f"✓ Key property found: {prop} = {business_props[prop]}")
            else:
                print(f"✗ Missing key property: {prop}")
                return False
                
    except Exception as e:
        print(f"✗ Error testing business properties: {e}")
        return False
    
    print("\n=== All Tests Passed! ===")
    return True

def test_visualization_integration():
    """Test visualization integration with different modes."""
    
    print("\n=== Testing Visualization Integration ===")
    
    try:
        node = ExactPostprocessNode()
        
        # Test each mode for visualization compatibility
        test_cases = [
            {
                'mode': 'fire_detection',
                'classes': 'No Fire,Fire',
                'expected_classes': 2,
                'description': 'Binary fire detection'
            },
            {
                'mode': 'yolo_v3', 
                'classes': 'person,car,bicycle,motorbike,bus',
                'expected_classes': 5,
                'description': 'Object detection'
            },
            {
                'mode': 'classification',
                'classes': 'cat,dog,bird,fish,rabbit',
                'expected_classes': 5,
                'description': 'Multi-class classification'
            }
        ]
        
        for case in test_cases:
            print(f"\n--- {case['description']} ---")
            
            # Configure node
            node.set_property('postprocess_type', case['mode'])
            node.set_property('class_names', case['classes'])
            
            # Get configuration for visualization
            config = node.get_postprocessing_config()
            parsed_classes = config['class_names']
            
            print(f"✓ Mode: {case['mode']}")
            print(f"✓ Classes: {parsed_classes}")
            print(f"✓ Expected {case['expected_classes']}, got {len(parsed_classes)}")
            
            if len(parsed_classes) == case['expected_classes']:
                print("✓ Class count matches expected")
            else:
                print(f"✗ Class count mismatch: expected {case['expected_classes']}, got {len(parsed_classes)}")
                return False
        
        print("\n✓ Visualization integration tests passed!")
        return True
        
    except Exception as e:
        print(f"✗ Error in visualization integration test: {e}")
        return False

if __name__ == "__main__":
    print("Starting ExactPostprocessNode Tests...\n")
    
    success = True
    
    # Run main functionality tests
    if not test_postprocess_node():
        success = False
    
    # Run visualization integration tests
    if not test_visualization_integration():
        success = False
    
    if success:
        print("\n🎉 All tests completed successfully!")
        print("ExactPostprocessNode is ready for mode switching and visualization!")
    else:
        print("\n❌ Some tests failed. Please check the implementation.")
        sys.exit(1)