cluster4npu/test_multi_series_fix.py

#!/usr/bin/env python3
"""
Test script to verify multi-series configuration fix
"""

import sys
import os
sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))

# Test the mflow_converter functionality
def test_multi_series_config_building():
    """Test building multi-series config from properties"""
    print("Testing multi-series config building...")
    
    from core.functions.mflow_converter import MFlowConverter
    
    # Create converter instance
    converter = MFlowConverter(default_fw_path='.')
    
    # Mock properties data that would come from a node
    test_properties = {
        'multi_series_mode': True,
        'enabled_series': ['520', '720'],
        'kl520_port_ids': '28,32',
        'kl720_port_ids': '4',
        'assets_folder': '',  # Empty for this test
        'max_queue_size': 100
    }
    
    # Test building config
    config = converter._build_multi_series_config_from_properties(test_properties)
    
    print(f"Generated config: {config}")
    
    if config:
        # Verify structure
        assert 'KL520' in config, "KL520 should be in config"
        assert 'KL720' in config, "KL720 should be in config"
        
        # Check KL520 config
        kl520_config = config['KL520']
        assert 'port_ids' in kl520_config, "KL520 should have port_ids"
        assert kl520_config['port_ids'] == [28, 32], f"KL520 port_ids should be [28, 32], got {kl520_config['port_ids']}"
        
        # Check KL720 config
        kl720_config = config['KL720']
        assert 'port_ids' in kl720_config, "KL720 should have port_ids"
        assert kl720_config['port_ids'] == [4], f"KL720 port_ids should be [4], got {kl720_config['port_ids']}"
        
        print("[OK] Multi-series config structure is correct")
    else:
        print("[ERROR] Config building returned None")
        return False
    
    # Test with invalid port IDs
    invalid_properties = {
        'multi_series_mode': True,
        'enabled_series': ['520'],
        'kl520_port_ids': 'invalid,port,ids',
        'assets_folder': ''
    }
    
    invalid_config = converter._build_multi_series_config_from_properties(invalid_properties)
    assert invalid_config is None, "Invalid port IDs should result in None config"
    print("[OK] Invalid port IDs handled correctly")
    
    return True

def test_stage_config():
    """Test StageConfig with multi-series support"""
    print("\\nTesting StageConfig with multi-series...")
    
    from core.functions.InferencePipeline import StageConfig
    
    # Test creating StageConfig with multi-series
    multi_series_config = {
        "KL520": {"port_ids": [28, 32]},
        "KL720": {"port_ids": [4]}
    }
    
    stage_config = StageConfig(
        stage_id="test_stage",
        port_ids=[],  # Not used in multi-series mode
        scpu_fw_path='',
        ncpu_fw_path='',
        model_path='',
        upload_fw=False,
        multi_series_mode=True,
        multi_series_config=multi_series_config
    )
    
    print(f"Created StageConfig with multi_series_mode: {stage_config.multi_series_mode}")
    print(f"Multi-series config: {stage_config.multi_series_config}")
    
    assert stage_config.multi_series_mode == True, "multi_series_mode should be True"
    assert stage_config.multi_series_config == multi_series_config, "multi_series_config should match"
    
    print("[OK] StageConfig supports multi-series configuration")
    return True

def main():
    """Run all tests"""
    print("Testing Multi-Series Configuration Fix")
    print("=" * 50)
    
    try:
        # Test config building
        if not test_multi_series_config_building():
            print("[ERROR] Config building test failed")
            return False
        
        # Test StageConfig
        if not test_stage_config():
            print("[ERROR] StageConfig test failed")
            return False
        
        print("\\n" + "=" * 50)
        print("[SUCCESS] All tests passed!")
        print("\\nThe fix should now properly:")
        print("1. Detect multi_series_mode from node properties")
        print("2. Build multi_series_config from series-specific port IDs")
        print("3. Pass the config to MultiDongle for true multi-series operation")
        
        return True
        
    except Exception as e:
        print(f"[ERROR] Test failed with exception: {e}")
        import traceback
        traceback.print_exc()
        return False

if __name__ == "__main__":
    success = main()
    sys.exit(0 if success else 1)