Add tests for wrapper