Fixed further bugs

- Now would be a good time to write tests