gh-124176: Add special support for dataclasses to `create_autospec` (#124429)

This commit is contained in:
sobolevn 2024-09-27 09:48:31 +03:00 committed by GitHub
parent 08e1bbe4a3
commit 3a0e7f5762
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
3 changed files with 92 additions and 6 deletions

View File

@ -8,8 +8,10 @@ from unittest.mock import (
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 @@ class SpecSignatureTest(unittest.TestCase):
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):

View File

@ -34,6 +34,7 @@ import builtins
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.

View File

@ -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.