led-controller/test/ws.py

import asyncio
import websockets
import json
import sys

async def test_websocket():
    uri = "ws://192.168.4.1:8080/ws"
    tests_passed = 0
    tests_total = 0
    
    async def run_test(name, test_func):
        nonlocal tests_passed, tests_total
        tests_total += 1
        try:
            result = await test_func()
            if result is not False:
                print(f"✓ {name}")
                tests_passed += 1
                return True
            else:
                print(f"✗ {name} (failed)")
                return False
        except Exception as e:
            print(f"✗ {name} (error: {e})")
            return False
    
    try:
        print(f"Connecting to WebSocket server at {uri}...")
        async with websockets.connect(uri) as websocket:
            print(f"✓ Connected to WebSocket server\n")
            
            # Test 1: Empty JSON
            print("Test 1: Empty JSON")
            await run_test("Send empty JSON", lambda: websocket.send(json.dumps({})))
            await asyncio.sleep(0.3)
            
            # Test 2: Pattern on with single color
            print("\nTest 2: Pattern 'on'")
            await run_test("Send on pattern", lambda: websocket.send(json.dumps({
                "settings": {"pattern": "on", "colors": ["#00ff00"], "brightness": 200}
            })))
            await asyncio.sleep(0.3)
            
            # Test 3: Pattern blink
            print("\nTest 3: Pattern 'blink'")
            await run_test("Send blink pattern", lambda: websocket.send(json.dumps({
                "settings": {"pattern": "blink", "colors": ["#ff0000"], "delay": 500}
            })))
            await asyncio.sleep(0.3)
            
            # Test 4: Pattern rainbow
            print("\nTest 4: Pattern 'rainbow'")
            await run_test("Send rainbow pattern", lambda: websocket.send(json.dumps({
                "settings": {"pattern": "rainbow", "delay": 100}
            })))
            await asyncio.sleep(0.3)
            
            # Test 5: Pattern off
            print("\nTest 5: Pattern 'off'")
            await run_test("Send off pattern", lambda: websocket.send(json.dumps({
                "settings": {"pattern": "off"}
            })))
            await asyncio.sleep(0.3)
            
            # Test 6: Multiple colors
            print("\nTest 6: Multiple colors")
            await run_test("Send multiple colors", lambda: websocket.send(json.dumps({
                "settings": {
                    "pattern": "color_transition",
                    "colors": ["#ff0000", "#00ff00", "#0000ff"],
                    "delay": 100
                }
            })))
            await asyncio.sleep(0.3)
            
            # Test 7: RGB tuple colors (if supported)
            print("\nTest 7: RGB tuple colors")
            await run_test("Send RGB tuple colors", lambda: websocket.send(json.dumps({
                "settings": {
                    "pattern": "on",
                    "colors": [[255, 0, 128], [128, 255, 0]],
                    "brightness": 150
                }
            })))
            await asyncio.sleep(0.3)
            
            # Test 8: Pattern with all parameters
            print("\nTest 8: Pattern with all parameters")
            await run_test("Send pattern with all params", lambda: websocket.send(json.dumps({
                "settings": {
                    "pattern": "flicker",
                    "colors": ["#ff8800"],
                    "brightness": 127,
                    "delay": 80,
                    "n1": 10,
                    "n2": 5,
                    "n3": 1,
                    "n4": 1
                }
            })))
            await asyncio.sleep(0.3)
            
            # Test 9: Short-key format (df/dj)
            print("\nTest 9: Short-key format (df/dj)")
            await run_test("Send df/dj format", lambda: websocket.send(json.dumps({
                "df": {"pt": "on", "cl": ["#ff0000"], "br": 200},
                "dj": {"pa": "blink", "cl": ["#00ff00"], "dl": 500},
                "settings": {"pattern": "blink", "colors": ["#00ff00"], "delay": 500, "brightness": 200}
            })))
            await asyncio.sleep(0.3)
            
            # Test 10: Rapid message sending
            print("\nTest 10: Rapid message sending")
            patterns = ["on", "off", "on", "blink"]
            for i, pattern in enumerate(patterns):
                p = pattern  # Capture in closure
                await run_test(f"Rapid send {i+1}/{len(patterns)}", lambda p=p: websocket.send(json.dumps({
                    "settings": {"pattern": p, "colors": ["#ffffff"]}
                })))
                await asyncio.sleep(0.1)
            
            # Test 11: Large message
            print("\nTest 11: Large message")
            large_colors = [f"#{i%256:02x}{i*2%256:02x}{i*3%256:02x}" for i in range(50)]
            await run_test("Send large message", lambda: websocket.send(json.dumps({
                "settings": {
                    "pattern": "color_transition",
                    "colors": large_colors,
                    "delay": 50
                }
            })))
            await asyncio.sleep(0.3)
            
            # Test 12: Invalid JSON (should be handled gracefully)
            print("\nTest 12: Invalid JSON handling")
            try:
                await websocket.send("not valid json")
                print("⚠ Invalid JSON sent (server should handle gracefully)")
                tests_total += 1
            except Exception as e:
                print(f"✗ Invalid JSON failed to send: {e}")
                tests_total += 1
            
            # Test 13: Malformed structure (missing settings)
            print("\nTest 13: Malformed structure")
            await run_test("Send message without settings", lambda: websocket.send(json.dumps({
                "pattern": "on",
                "colors": ["#ff0000"]
            })))
            await asyncio.sleep(0.3)
            
            # Test 14: Just settings key, no pattern
            print("\nTest 14: Settings without pattern")
            await run_test("Send settings without pattern", lambda: websocket.send(json.dumps({
                "settings": {"colors": ["#0000ff"], "brightness": 100}
            })))
            await asyncio.sleep(0.3)
            
            # Test 15: Empty settings
            print("\nTest 15: Empty settings")
            await run_test("Send empty settings", lambda: websocket.send(json.dumps({
                "settings": {}
            })))
            await asyncio.sleep(0.3)
            
            print(f"\n{'='*50}")
            print(f"Tests completed: {tests_passed}/{tests_total} passed")
            if tests_passed == tests_total:
                print("✓ All tests passed!")
            else:
                print(f"⚠ {tests_total - tests_passed} test(s) failed")
            print(f"{'='*50}")
            
    except websockets.exceptions.ConnectionClosedOK:
        print("✓ WebSocket connection closed gracefully.")
    except websockets.exceptions.ConnectionClosedError as e:
        print(f"✗ WebSocket connection closed with error: {e}")
        sys.exit(1)
    except ConnectionRefusedError:
        print(f"✗ Connection refused. Is the server running at {uri}?")
        print("Make sure:")
        print("  1. The device is connected to WiFi")
        print("  2. The server is running on the device")
        print("  3. You can reach 192.168.4.1")
        sys.exit(1)
    except Exception as e:
        print(f"✗ An unexpected error occurred: {e}")
        import traceback
        traceback.print_exc()
        sys.exit(1)

if __name__ == "__main__":
    asyncio.run(test_websocket())