test_collection_connection.py

"""
Quick test to verify Arduino is sending data in the right format
"""

import serial
import serial.tools.list_ports
import time

def test_collection_format(port='COM3', baud_rate=9600, duration=10):
    """Test if Arduino is outputting data in the format expected by collect_episodes.py"""
    try:
        print(f"Connecting to {port} at {baud_rate} baud...")
        ser = serial.Serial(port, baud_rate, timeout=1)
        print("✓ Connected!\n")
        print("Waiting for data (10 seconds)...")
        print("=" * 60)
        
        start_time = time.time()
        line_count = 0
        has_features = False
        has_initialization = False
        
        while time.time() - start_time < duration:
            if ser.in_waiting:
                line = ser.readline().decode('utf-8', errors='ignore').strip()
                if line:
                    line_count += 1
                    print(f"[{line_count}] {line[:100]}")
                    
                    # Check for expected formats
                    if 'Features:' in line:
                        has_features = True
                        print("  ✓ Found 'Features:' - Good! This is what collect_episodes.py expects")
                    if 'Initialized' in line or 'Ready' in line:
                        has_initialization = True
                        print("  ✓ Found initialization message")
            else:
                time.sleep(0.1)
        
        print("=" * 60)
        print(f"\nSummary:")
        print(f"  Total lines received: {line_count}")
        print(f"  Has 'Features:' format: {'✓ YES' if has_features else '✗ NO'}")
        print(f"  Has initialization: {'✓ YES' if has_initialization else '✗ NO'}")
        
        if not has_features:
            print("\n⚠ WARNING: No 'Features:' lines found!")
            print("  The collect_episodes.py script expects data with 'Features:' in it.")
            print("  Make sure you uploaded arduino_slip_detection.ino (not arduino_raw_sensor.ino)")
            print("  The Arduino should output lines like:")
            print("  'Features: [0.123, 0.456, ...] | Prediction: NO_SLIP | Probability: 0.123'")
        else:
            print("\n✓ Arduino is outputting data in the correct format!")
            print("  You can now run: python collect_episodes.py --port COM3")
        
        ser.close()
        
    except serial.SerialException as e:
        print(f"✗ Error: {e}")
        print("\nAvailable ports:")
        for p in serial.tools.list_ports.comports():
            print(f"  - {p.device}: {p.description}")

if __name__ == '__main__':
    import argparse
    parser = argparse.ArgumentParser(description='Test Arduino output format for collection')
    parser.add_argument('--port', type=str, default='COM3', help='Serial port')
    parser.add_argument('--baud', type=int, default=9600, help='Baud rate')
    parser.add_argument('--duration', type=int, default=10, help='Test duration in seconds')
    
    args = parser.parse_args()
    test_collection_format(args.port, args.baud, args.duration)