Fix unit tests with incorrect mocking