Fixed issues from review and added TODO for missing test cases