go fmt
Refactor mock using polymorphism for cleaner code