Fixed a couple issues, added a test