Adding back in code that should fix unit test fails