Add in wait to test where it's currently failing.