Fix hang bug

Fix bug that hangs unit tests.