Fixed a few issues with tests