Catch exceptions in test functions.

Also, refactor the code quite a lot.