Cleaned up existing unit tests