style(mock): code cleanup

Closes #597