Updated for refactored framework