Fix broken tests

Minor changes to fix two broken tests