Replaced a bit of code to fix the unit tests.