Refactor to use injections