Fixed fallout from refactoring