Fixing unit tests to mock correctly