Forgot a little something
Slight changes to make the unit tests work again