Completed and wrote unit tests for #4