diff --git a/stdlib/@tests/test_cases/check_unittest.py b/stdlib/@tests/test_cases/check_unittest.py index 40c6efaa8ca0..d717e99a3f8d 100644 --- a/stdlib/@tests/test_cases/check_unittest.py +++ b/stdlib/@tests/test_cases/check_unittest.py @@ -5,9 +5,9 @@ from datetime import datetime, timedelta from decimal import Decimal from fractions import Fraction -from typing import TypedDict +from typing import TypedDict, Union from typing_extensions import assert_type -from unittest.mock import MagicMock, Mock, patch +from unittest.mock import AsyncMock, MagicMock, Mock, patch case = unittest.TestCase() @@ -154,10 +154,17 @@ def f_explicit_new(i: int) -> str: return "asdf" +@patch("sys.exit", new_callable=lambda: 42) +def f_explicit_new_callable(i: int, new_callable_ret: int) -> str: + return "asdf" + + assert_type(f_default_new(1), str) f_default_new("a") # Not an error due to ParamSpec limitations assert_type(f_explicit_new(1), str) f_explicit_new("a") # type: ignore[arg-type] +assert_type(f_explicit_new_callable(1), str) +f_explicit_new_callable("a") # Same as default new @patch("sys.exit", new=Mock()) @@ -171,3 +178,51 @@ def method() -> int: assert_type(TestXYZ.attr, int) assert_type(TestXYZ.method(), int) + + +with patch("sys.exit") as default_new_enter: + assert_type(default_new_enter, Union[MagicMock, AsyncMock]) + +with patch("sys.exit", new=42) as explicit_new_enter: + assert_type(explicit_new_enter, int) + +with patch("sys.exit", new_callable=lambda: 42) as explicit_new_callable_enter: + assert_type(explicit_new_callable_enter, int) + + +### +# Tests for mock.patch.object +### + + +@patch.object(Decimal, "exp") +def obj_f_default_new(i: int, mock: MagicMock) -> str: + return "asdf" + + +@patch.object(Decimal, "exp", new=42) +def obj_f_explicit_new(i: int) -> str: + return "asdf" + + +@patch.object(Decimal, "exp", new_callable=lambda: 42) +def obj_f_explicit_new_callable(i: int, new_callable_ret: int) -> str: + return "asdf" + + +assert_type(obj_f_default_new(1), str) +obj_f_default_new("a") # Not an error due to ParamSpec limitations +assert_type(obj_f_explicit_new(1), str) +obj_f_explicit_new("a") # type: ignore[arg-type] +assert_type(obj_f_explicit_new_callable(1), str) +obj_f_explicit_new_callable("a") # Same as default new + + +with patch.object(Decimal, "exp") as obj_default_new_enter: + assert_type(obj_default_new_enter, Union[MagicMock, AsyncMock]) + +with patch.object(Decimal, "exp", new=42) as obj_explicit_new_enter: + assert_type(obj_explicit_new_enter, int) + +with patch.object(Decimal, "exp", new_callable=lambda: 42) as obj_explicit_new_callable_enter: + assert_type(obj_explicit_new_callable_enter, int) diff --git a/stdlib/unittest/mock.pyi b/stdlib/unittest/mock.pyi index 9e353900f2d7..ec5fced32ea4 100644 --- a/stdlib/unittest/mock.pyi +++ b/stdlib/unittest/mock.pyi @@ -262,7 +262,7 @@ class _patch(Generic[_T]): # This class does not exist at runtime, it's a hack to make this work: # @patch("foo") # def bar(..., mock: MagicMock) -> None: ... -class _patch_default_new(_patch[MagicMock | AsyncMock]): +class _patch_pass_arg(_patch[_T]): @overload def __call__(self, func: _TT) -> _TT: ... # Can't use the following as ParamSpec is only allowed as last parameter: @@ -303,7 +303,7 @@ class _patcher: create: bool = ..., spec_set: Any | None = ..., autospec: Any | None = ..., - new_callable: Any | None = ..., + new_callable: Callable[..., Any] | None = ..., **kwargs: Any, ) -> _patch[_T]: ... @overload @@ -315,9 +315,21 @@ class _patcher: create: bool = ..., spec_set: Any | None = ..., autospec: Any | None = ..., - new_callable: Any | None = ..., + new_callable: Callable[..., _T], **kwargs: Any, - ) -> _patch_default_new: ... + ) -> _patch_pass_arg[_T]: ... + @overload + def __call__( + self, + target: str, + *, + spec: Any | None = ..., + create: bool = ..., + spec_set: Any | None = ..., + autospec: Any | None = ..., + new_callable: None = ..., + **kwargs: Any, + ) -> _patch_pass_arg[MagicMock | AsyncMock]: ... @overload @staticmethod def object( @@ -328,7 +340,7 @@ class _patcher: create: bool = ..., spec_set: Any | None = ..., autospec: Any | None = ..., - new_callable: Any | None = ..., + new_callable: Callable[..., Any] | None = ..., **kwargs: Any, ) -> _patch[_T]: ... @overload @@ -341,9 +353,22 @@ class _patcher: create: bool = ..., spec_set: Any | None = ..., autospec: Any | None = ..., - new_callable: Any | None = ..., + new_callable: Callable[..., _T], + **kwargs: Any, + ) -> _patch_pass_arg[_T]: ... + @overload + @staticmethod + def object( + target: Any, + attribute: str, + *, + spec: Any | None = ..., + create: bool = ..., + spec_set: Any | None = ..., + autospec: Any | None = ..., + new_callable: None = ..., **kwargs: Any, - ) -> _patch[MagicMock | AsyncMock]: ... + ) -> _patch_pass_arg[MagicMock | AsyncMock]: ... @staticmethod def multiple( target: Any,