Fix a race condition when "expected_state" is used

Fixed a race condition when comparing the "expected_state" while the new state wasn't reflected yet.