Fix unit test to account for device backends
Refactor common code into device.py