Fix stalling token unit tests
Clean up the token unit tests a bit