Refactored production code to reflect actual classes