Changes to TODO unit test