Refactored code, updated/simplified tests