Moved mock class from production code to test