Fix failing tests - good old red, green, refactor :-)