Adding test cases and fixing code to work with tests