Fix the tests failing, happy day!