Skip to content

Commit

Permalink
gh-124176: Add special support for dataclasses to create_autospec (#…
Browse files Browse the repository at this point in the history
  • Loading branch information
sobolevn authored Sep 27, 2024
1 parent 08e1bbe commit 3a0e7f5
Show file tree
Hide file tree
Showing 3 changed files with 92 additions and 6 deletions.
72 changes: 72 additions & 0 deletions Lib/test/test_unittest/testmock/testhelpers.py
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down Expand Up @@ -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):

Expand Down
22 changes: 16 additions & 6 deletions Lib/unittest/mock.py
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down Expand Up @@ -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:
Expand Down Expand Up @@ -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
Expand All @@ -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.
Expand Down
Original file line number Diff line number Diff line change
@@ -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.

0 comments on commit 3a0e7f5

Please sign in to comment.