Switched to use mockery for testing