Refactored code to pass all current integration tests