Updated with better tests