Fix a small bug in function wrapping and add some more tests