Add hook for test framework instead of relying on custom code. Fixes #2