#!/usr/bin/env python

import usb
import sys
from time import sleep

# For control transfer test
CTRL_LOOPBACK_WRITE = 0
CTRL_LOOPBACK_READ = 1

def ctrl_write(handle, data, timeout = 1000):
	written = handle.controlMsg(0x40, CTRL_LOOPBACK_WRITE, data, 0, 0, timeout)
	assert written == len(data)

def ctrl_read(handle, size, timeout = 1000):
	return handle.controlMsg(0xC0, CTRL_LOOPBACK_READ, size, 0, 0, timeout)

# Find the specified device.
# If it is not found, return None
def find_device(busses, idProduct, idVendor):
	for bus in busses:
		for dev in bus.devices:
			if dev.idProduct == idProduct and dev.idVendor == idVendor:
				return dev


# Loopback test for bulk endpoint
def test_bulk(handle, msg):
	handle.bulkWrite(0x1, msg, 1000)
	data = handle.bulkRead(0x81, len(msg), 1000)
	str = "".join([chr(i) for i in data]) # transforma em string
	assert str == msg

# Loopback test for interrupt endpoint
def test_interrupt(handle, msg):
	handle.interruptWrite(0x2, msg, 1000)
	data = handle.interruptRead(0x82, len(msg), 1000)
	str = "".join([chr(i) for i in data])
	assert str == msg
	
# Loopback test for control transfers
def test_control(handle, msg):
	ctrl_write(handle, msg, 1000)
	data = ctrl_read(handle,  len(msg), 1000)
	str = "".join([chr(i) for i in data])
	assert len(str) == len(msg)
	assert str == msg


if __name__ == "__main__":
	print "********************************"
	print "PyUSB test script"
	print "********************************"
	print ""

	busses = usb.busses()

	print "enumeration device test..."
	dev = find_device(busses, 0x0001, 0xFFFE)
	if dev is None:
		print "enumeration test failed..."
		sys.exit(1)
	
	print "enumeration test ok..."
	
	print "device open test..."
	handle = dev.open()
	print "device open ok..."

	print "configuration test..."
	handle.setConfiguration(1)
	print "configuration test ok..."

	print "claim interface test..."
	handle.claimInterface(0)
	print "claim interface test ok..."

	print "interface configuration test..."
	handle.setAltInterface(0)
	print "interface configuration ok..."

	print "I/O test..."

	for x in range(10):
		test_interrupt(handle, "interrupt test 1")
		test_bulk(handle, "bulk test 1")
		test_control(handle, "control test 1")
		test_interrupt(handle, "interrupt test 2")
		test_bulk(handle, "bulk test 2")
		test_control(handle, "control test 2")

	print "I/O test ok..."

	print "reset endpoint test..."
	handle.clearHalt(1)	
	print "reset endpoint ok..."

	print "string test..."
	str1 = handle.getString(1,100)
	str2 = handle.getString(2,100)
	print "string index 1:", str1
	print "string index 2:", str2
	assert str1 == 'Mxyzp7lk'
	assert str2 == 'PyUSB'
	print "string test ok..."

	# Get the device descriptor to check if getDescriptor works fine
	print "descriptor test..."
	descr = handle.getDescriptor(1, 0, 18)
	assert len(descr) == 18
	print "descriptor:", descr
	print "descriptor test ok..."

	# Do a GET_STATUS control request to check if control transfer is ok
	print "control transfer test..."
	control_res = handle.controlMsg(0x80, 0, 2, 0, 0, 1000)
	assert len(control_res) == 2
	print "control tranfer result: ", control_res
	print "control transfer ok..."

	print "reset device test..."
	handle.reset()

	del handle
	del dev
	del busses

	sleep(3) # Necessary to Work in Windows

	busses = usb.busses()

	dev = find_device(busses, 0x0001, 0xFFFE)
	assert dev is not None

	del dev
	del busses
	del usb
	
	print "type enter key to exit"
	raw_input()
