Attempted fix to the tests