Fix tests

@localshred two tests here I didn't know how to fix. Take a look?