Add missed class for previous commit