Fix a potential hang in unit tests