Refactor check_for_state/wait_for_state

Signed-off-by: Maxime Petazzoni <max@signalfuse.com>