Improve async usage in tests.