Updated virtual_machines.py

* Use ``wait_for()`` to be faster when checking the vm state change
* Be nice to PEP8