Add TODO in comments for unit test