From 4f9fe23a03afc5c356d40e43c94b1d86797d264b Mon Sep 17 00:00:00 2001 From: Richard Ross Date: Mon, 19 Dec 2022 14:23:03 -0800 Subject: [PATCH] Fix default_context_manager in conjunction with ExitStack. (#344) Summary: Pull Request resolved: https://github.com/facebook/TestSlide/pull/344 As outlined in the comment; because our MethodProxy objects don't get bound to the object (via __get__); when (Async)ExitStack comes along and invokes: ``` enter_fn = cm.__class_.__enter__ enter_fn(cm) ``` A problem occurs - while the attribute lookup succeeds (yay); it's not expecting a `self` paramter explicitly passed here, as it's not a bound python method. Theoretically; magic could be added to MethodProxy to properly bind via `__get__`, but my knowledge of testslide right now isn't strong enough to feel confident in that approach, so I'm opting for this. Reviewed By: deathowl Differential Revision: D42114206 fbshipit-source-id: 9473fe59d30c93c56717f92117180f1b6f9e31c1 --- tests/strict_mock_testslide.py | 16 +++++++ testslide/strict_mock.py | 77 +++++++++++++++++++++++++++++----- 2 files changed, 82 insertions(+), 11 deletions(-) diff --git a/tests/strict_mock_testslide.py b/tests/strict_mock_testslide.py index 27a8a9f..084a9de 100644 --- a/tests/strict_mock_testslide.py +++ b/tests/strict_mock_testslide.py @@ -814,6 +814,14 @@ def it_yields_the_mock(self): with self.strict_mock as target: self.assertTrue(target is self.strict_mock) + @context.example + def works_with_exitstack(self): + with contextlib.ExitStack() as exit_stack: + target = exit_stack.enter_context( + self.strict_mock + ) + self.assertTrue(target is self.strict_mock) + @context.sub_context def string_template(context): @context.example @@ -1116,6 +1124,14 @@ async def it_yields_the_mock(self): async with self.strict_mock as m: assert id(self.strict_mock) == id(m) + @context.example + async def works_with_exitstack(self): + async with contextlib.AsyncExitStack() as exit_stack: + target = await exit_stack.enter_async_context( + self.strict_mock + ) + self.assertTrue(target is self.strict_mock) + @context.sub_context def making_copies(context): context.memoize("strict_mock", lambda self: StrictMock(template=Template)) diff --git a/testslide/strict_mock.py b/testslide/strict_mock.py index f70a0cb..49def41 100644 --- a/testslide/strict_mock.py +++ b/testslide/strict_mock.py @@ -430,25 +430,80 @@ def __setup_magic_methods(self) -> None: ): setattr(self, name, _DefaultMagic(self, name)) + # THE PROBLEM: + # Certain idiomatic python use-cases (*cough* (Async)ExitStack *cough*) + # invoke context managers with roughly the following code: + # + # enter_fn = cm.__class__.__aenter__ + # await enter_fn(cm) + # + # While fine for 'normal' python objects, this is problematic for how our + # MethodProxy builds functions; as they don't end up properly bound + # (i.e. via __get__) to the object; and end up failing with cryptic errors + # like: + # + # 'aenter invoked with 1 arguments; expected 0'. + # + # To fix this; let's properly define these as real method on the object so + # they get bound. They still will not be visible to outside code due to our + # validating in __getattr__ that they match relevant entries in the template + def __enter_default_context_manager_helper__(self): + return self + + def __exit_default_context_manager_helper__(self, exc_type, exc_value, traceback): + pass + + async def __aenter_default_context_manager_helper__(self): + return self + + async def __aexit_default_context_manager_helper__( + self, exc_type, exc_value, traceback + ): + pass + def __setup_default_context_manager(self, default_context_manager: bool) -> None: if self._template and default_context_manager: if hasattr(self._template, "__enter__") and hasattr( self._template, "__exit__" ): - self.__enter__ = lambda: self - self.__exit__ = lambda exc_type, exc_value, traceback: None + # Use setattr(type(self)) to explicitly avoid any MethodProxy + # magic which breaks the method-bind. + # Also needed for getattr to bypass our template class check. + setattr( # noqa: B010 + type(self), + "__enter__", + getattr( # noqa: B009 + type(self), "__enter_default_context_manager_helper__" + ), + ) + setattr( # noqa: B010 + type(self), + "__exit__", + getattr( # noqa: B009 + type(self), "__exit_default_context_manager_helper__" + ), + ) + if hasattr(self._template, "__aenter__") and hasattr( self._template, "__aexit__" ): - - async def aenter(): - return self - - async def aexit(exc_type, exc_value, traceback): - pass - - self.__aenter__ = aenter - self.__aexit__ = aexit + # Use setattr(type(self)) to explicitly avoid any MethodProxy + # magic which breaks the method-bind. + # Also needed for getattr to bypass our template class check. + setattr( # noqa: B010 + type(self), + "__aenter__", + getattr( # noqa: B009 + type(self), "__aenter_default_context_manager_helper__" + ), + ) + setattr( # noqa: B010 + type(self), + "__aexit__", + getattr( # noqa: B009 + type(self), "__aexit_default_context_manager_helper__" + ), + ) def __get_caller_frame(self, depth: int) -> FrameType: # Adding extra 3 to account for the stack: