Made some changes to the unit tests due to my next blog post