Added unit tests for problem 2