Rework wrap callback. Add test to improve coverage