Cleaned up existing unit tests, added some new ones!