StreamLens/test_button_layout_improvements.py

#!/usr/bin/env python3
"""
Test script to verify the button layout and table sorting improvements
"""

import sys
import time
from pathlib import Path

# Add analyzer to path
sys.path.insert(0, str(Path(__file__).parent))

from analyzer.analysis.core import EthernetAnalyzer
from analyzer.tui.textual.widgets.filtered_flow_view import FilteredFlowView


def test_button_layout():
    """Test that buttons are now one row high"""
    print("Testing button layout improvements...")
    
    # Create analyzer and flow view
    analyzer = EthernetAnalyzer()
    flow_view = FilteredFlowView(analyzer)
    
    # Check that button height is set to 1
    css_content = flow_view.DEFAULT_CSS
    print("Checking CSS for button height...")
    
    if "height: 1;" in css_content:
        print("✅ Button height set to 1 row")
    else:
        print("❌ Button height not set correctly")
        return False
    
    if "#filter-bar {\n        height: 1;" in css_content:
        print("✅ Filter bar height set to 1 row")
    else:
        print("❌ Filter bar height not set correctly")
        return False
    
    return True


def test_button_ordering():
    """Test that buttons will be ordered by count"""
    print("\nTesting button ordering by count...")
    
    analyzer = EthernetAnalyzer()
    
    # Add mock flow data with different frame type counts
    from analyzer.models.flow_stats import FlowStats, FrameTypeStats
    
    # Flow 1: High CH10-Data count
    flow1 = FlowStats(src_ip="192.168.1.1", dst_ip="192.168.1.2")
    flow1.frame_types["CH10-Data"] = FrameTypeStats("CH10-Data", count=1000)
    flow1.frame_types["UDP"] = FrameTypeStats("UDP", count=10)
    analyzer.flows["flow1"] = flow1
    
    # Flow 2: High UDP count
    flow2 = FlowStats(src_ip="192.168.1.3", dst_ip="192.168.1.4")
    flow2.frame_types["UDP"] = FrameTypeStats("UDP", count=500)
    flow2.frame_types["PTP-Sync"] = FrameTypeStats("PTP-Sync", count=200)
    flow2.frame_types["TMATS"] = FrameTypeStats("TMATS", count=5)
    analyzer.flows["flow2"] = flow2
    
    flow_view = FilteredFlowView(analyzer)
    
    # Get frame types and their counts
    frame_types = flow_view._get_all_frame_types()
    print(f"Frame types detected: {frame_types}")
    
    # Check that they're sorted by count (CH10-Data: 1000, UDP: 510, PTP-Sync: 200, TMATS: 5)
    sorted_types = sorted(frame_types.items(), key=lambda x: x[1], reverse=True)
    print(f"Sorted frame types: {sorted_types}")
    
    if sorted_types[0][0] == "CH10-Data" and sorted_types[0][1] == 1000:
        print("✅ Highest count frame type (CH10-Data) will be first")
    else:
        print(f"❌ Expected CH10-Data (1000) first, got {sorted_types[0]}")
        return False
    
    if sorted_types[1][0] == "UDP" and sorted_types[1][1] == 510:
        print("✅ Second highest count frame type (UDP) will be second")
    else:
        print(f"❌ Expected UDP (510) second, got {sorted_types[1]}")
        return False
    
    return True


def test_table_sorting():
    """Test table sorting functionality"""
    print("\nTesting table sorting functionality...")
    
    analyzer = EthernetAnalyzer()
    flow_view = FilteredFlowView(analyzer)
    
    # Check that sorting state is initialized
    if hasattr(flow_view, 'sort_column') and hasattr(flow_view, 'sort_reverse'):
        print("✅ Sorting state variables initialized")
    else:
        print("❌ Sorting state variables not found")
        return False
    
    # Check that sort method exists
    if hasattr(flow_view, 'action_sort_column'):
        print("✅ Sort action method exists")
    else:
        print("❌ Sort action method not found")
        return False
    
    # Check that get_sort_key method exists
    if hasattr(flow_view, '_get_sort_key'):
        print("✅ Sort key method exists")
    else:
        print("❌ Sort key method not found")
        return False
    
    # Test sort key extraction
    test_row = ["1", "192.168.1.1:5000", "239.1.1.1:8000", "UDP", "1,234", "Normal"]
    
    # Test numeric extraction
    key = flow_view._get_sort_key(test_row, 4)  # Packet count column
    if key == 1234:
        print("✅ Numeric sort key extraction works (comma removal)")
    else:
        print(f"❌ Expected 1234, got {key}")
        return False
    
    # Test string extraction
    key = flow_view._get_sort_key(test_row, 3)  # Protocol column
    if key == "udp":
        print("✅ String sort key extraction works (lowercase)")
    else:
        print(f"❌ Expected 'udp', got {key}")
        return False
    
    return True


def test_key_bindings():
    """Test that key bindings are set up correctly"""
    print("\nTesting key bindings...")
    
    # Import the app to check key bindings
    from analyzer.tui.textual.app_v2 import StreamLensAppV2
    from analyzer.tui.textual.widgets.filtered_flow_view import FilteredFlowView
    
    # Check FilteredFlowView bindings
    flow_view_bindings = [binding.key for binding in FilteredFlowView.BINDINGS]
    expected_bindings = ['alt+1', 'alt+2', 'alt+3', 'alt+4', 'alt+5', 'alt+6', 'alt+7', 'alt+8', 'alt+9', 'alt+0']
    
    for binding in expected_bindings:
        if binding in flow_view_bindings:
            print(f"✅ {binding} binding found in FilteredFlowView")
        else:
            print(f"❌ {binding} binding missing in FilteredFlowView")
            return False
    
    # Check main app bindings
    app_bindings = [binding[0] for binding in StreamLensAppV2.BINDINGS]
    
    for binding in expected_bindings:
        if binding in app_bindings:
            print(f"✅ {binding} binding found in main app")
        else:
            print(f"❌ {binding} binding missing in main app")
            return False
    
    return True


if __name__ == "__main__":
    print("StreamLens Button Layout & Sorting Test")
    print("=" * 50)
    
    try:
        success1 = test_button_layout()
        success2 = test_button_ordering()
        success3 = test_table_sorting()
        success4 = test_key_bindings()
        
        if success1 and success2 and success3 and success4:
            print(f"\n✅ All tests passed!")
            print(f"\n📊 Summary of Improvements:")
            print(f"   • Buttons are now 1 row high (was 3)")
            print(f"   • Buttons ordered by frame type count (highest first)")
            print(f"   • Tables sortable with Alt+1...Alt+0 keys")
            print(f"   • Smart sorting handles numbers, text, and units")
            print(f"\n🎯 Usage:")
            print(f"   • 1-9,0: Select frame type filters")
            print(f"   • Alt+1...Alt+0: Sort by columns 1-10")
            print(f"   • Alt+same key: Toggle sort direction")
        else:
            print(f"\n❌ Some tests failed")
            sys.exit(1)
            
    except Exception as e:
        print(f"\n❌ Test failed with error: {e}")
        import traceback
        traceback.print_exc()
        sys.exit(1)