Fix test wait