Refactor test to stub appropriate object