Bugfix - Test cleanup

This PR mostly incorporates the work done by @RafeKettler in #183.

I've just cleaned up some details.

Thanks for the contribution @RafeKettler! :+1: