Code Clean UP + Tests