From 3a0e7f57628466aedcaaf6c5ff7c8224f5155a2c Mon Sep 17 00:00:00 2001 From: sobolevn Date: Fri, 27 Sep 2024 09:48:31 +0300 Subject: [PATCH] gh-124176: Add special support for dataclasses to `create_autospec` (#124429) --- .../test_unittest/testmock/testhelpers.py | 72 +++++++++++++++++++ Lib/unittest/mock.py | 22 ++++-- ...-09-24-13-32-16.gh-issue-124176.6hmOPz.rst | 4 ++ 3 files changed, 92 insertions(+), 6 deletions(-) create mode 100644 Misc/NEWS.d/next/Library/2024-09-24-13-32-16.gh-issue-124176.6hmOPz.rst diff --git a/Lib/test/test_unittest/testmock/testhelpers.py b/Lib/test/test_unittest/testmock/testhelpers.py index c9c20f008ca5a2..f260769eb8c35e 100644 --- a/Lib/test/test_unittest/testmock/testhelpers.py +++ b/Lib/test/test_unittest/testmock/testhelpers.py @@ -8,8 +8,10 @@ Mock, ANY, _CallList, patch, PropertyMock, _callable ) +from dataclasses import dataclass, field, InitVar from datetime import datetime from functools import partial +from typing import ClassVar class SomeClass(object): def one(self, a, b): pass @@ -1034,6 +1036,76 @@ def f(a): pass self.assertEqual(mock.mock_calls, []) self.assertEqual(rv.mock_calls, []) + def test_dataclass_post_init(self): + @dataclass + class WithPostInit: + a: int = field(init=False) + b: int = field(init=False) + def __post_init__(self): + self.a = 1 + self.b = 2 + + for mock in [ + create_autospec(WithPostInit, instance=True), + create_autospec(WithPostInit()), + ]: + with self.subTest(mock=mock): + self.assertIsInstance(mock.a, int) + self.assertIsInstance(mock.b, int) + + # Classes do not have these fields: + mock = create_autospec(WithPostInit) + msg = "Mock object has no attribute" + with self.assertRaisesRegex(AttributeError, msg): + mock.a + with self.assertRaisesRegex(AttributeError, msg): + mock.b + + def test_dataclass_default(self): + @dataclass + class WithDefault: + a: int + b: int = 0 + + for mock in [ + create_autospec(WithDefault, instance=True), + create_autospec(WithDefault(1)), + ]: + with self.subTest(mock=mock): + self.assertIsInstance(mock.a, int) + self.assertIsInstance(mock.b, int) + + def test_dataclass_with_method(self): + @dataclass + class WithMethod: + a: int + def b(self) -> int: + return 1 + + for mock in [ + create_autospec(WithMethod, instance=True), + create_autospec(WithMethod(1)), + ]: + with self.subTest(mock=mock): + self.assertIsInstance(mock.a, int) + mock.b.assert_not_called() + + def test_dataclass_with_non_fields(self): + @dataclass + class WithNonFields: + a: ClassVar[int] + b: InitVar[int] + + msg = "Mock object has no attribute" + for mock in [ + create_autospec(WithNonFields, instance=True), + create_autospec(WithNonFields(1)), + ]: + with self.subTest(mock=mock): + with self.assertRaisesRegex(AttributeError, msg): + mock.a + with self.assertRaisesRegex(AttributeError, msg): + mock.b class TestCallList(unittest.TestCase): diff --git a/Lib/unittest/mock.py b/Lib/unittest/mock.py index df3901f9660ac1..21ca061a77c26f 100644 --- a/Lib/unittest/mock.py +++ b/Lib/unittest/mock.py @@ -34,6 +34,7 @@ import pkgutil from inspect import iscoroutinefunction import threading +from dataclasses import fields, is_dataclass from types import CodeType, ModuleType, MethodType from unittest.util import safe_repr from functools import wraps, partial @@ -2756,7 +2757,15 @@ def create_autospec(spec, spec_set=False, instance=False, _parent=None, raise InvalidSpecError(f'Cannot autospec a Mock object. ' f'[object={spec!r}]') is_async_func = _is_async_func(spec) - _kwargs = {'spec': spec} + + entries = [(entry, _missing) for entry in dir(spec)] + if is_type and instance and is_dataclass(spec): + dataclass_fields = fields(spec) + entries.extend((f.name, f.type) for f in dataclass_fields) + _kwargs = {'spec': [f.name for f in dataclass_fields]} + else: + _kwargs = {'spec': spec} + if spec_set: _kwargs = {'spec_set': spec} elif spec is None: @@ -2813,7 +2822,7 @@ def create_autospec(spec, spec_set=False, instance=False, _parent=None, _name='()', _parent=mock, wraps=wrapped) - for entry in dir(spec): + for entry, original in entries: if _is_magic(entry): # MagicMock already does the useful magic methods for us continue @@ -2827,10 +2836,11 @@ def create_autospec(spec, spec_set=False, instance=False, _parent=None, # AttributeError on being fetched? # we could be resilient against it, or catch and propagate the # exception when the attribute is fetched from the mock - try: - original = getattr(spec, entry) - except AttributeError: - continue + if original is _missing: + try: + original = getattr(spec, entry) + except AttributeError: + continue child_kwargs = {'spec': original} # Wrap child attributes also. diff --git a/Misc/NEWS.d/next/Library/2024-09-24-13-32-16.gh-issue-124176.6hmOPz.rst b/Misc/NEWS.d/next/Library/2024-09-24-13-32-16.gh-issue-124176.6hmOPz.rst new file mode 100644 index 00000000000000..38c030668b6b42 --- /dev/null +++ b/Misc/NEWS.d/next/Library/2024-09-24-13-32-16.gh-issue-124176.6hmOPz.rst @@ -0,0 +1,4 @@ +Add support for :func:`dataclasses.dataclass` in +:func:`unittest.mock.create_autospec`. Now ``create_autospec`` will check +for potential dataclasses and use :func:`dataclasses.fields` function to +retrieve the spec information.