add tests for problem 4