"""
Test that our implementation of wrap_non_picklable_objects mimics
properly the loky implementation.
"""

from .._cloudpickle_wrapper import (
    _my_wrap_non_picklable_objects,
    wrap_non_picklable_objects,
)


def a_function(x):
    return x


class AClass(object):
    def __call__(self, x):
        return x


def test_wrap_non_picklable_objects():
    # Mostly a smoke test: test that we can use callable in the same way
    # with both our implementation of wrap_non_picklable_objects and the
    # upstream one
    for obj in (a_function, AClass()):
        wrapped_obj = wrap_non_picklable_objects(obj)
        my_wrapped_obj = _my_wrap_non_picklable_objects(obj)
        assert wrapped_obj(1) == my_wrapped_obj(1)