Fix incorrect assertion.
Generalize device state context manager.