Updated tests

Just some small tweaks to the tests