CODEWITHSUNDEEP
pythonattributesmockingpython-dataclasses
mohi666
mohi666

Reputation: 7052

How to use spec when mocking data classes in Python

I'm trying to port our namedtuple classes into dataclass in Python 3.6 using the backport package. However, I noticed when mocking dataclass classes, you cannot use the "spec" keyword anymore. I assume it's because the dataclass code is auto generated.

from dataclasses import dataclass
import mock

@dataclass
class A:
    aaa: str
    bbb: int


m = mock.Mock(spec=A)

m.aaa

And this is the error I get:

AttributeError: Mock object has no attribute 'aaa'

Any idea if there's any way to automatically set all the attributes from original object to the mock object? I have lots of data classes with a lot of data. It's going to be really tedious if I try to manually set the values one by one.

Upvotes: 15

Views: 9322

Answers (4)

cowlinator
cowlinator

Reputation: 8813

Dataclass fields are actually implemented as instance variables. You have 4 options:

  • give the field a default value:
@dataclass()
class MyDataClass1:
    my_field: str = field(default="my_field_val")
m = Mock(spec=MyDataClass1("my_field_val"))
print(m.my_field)
  • provide the value for the field as an arg in the constructor call of the dataclass:
@dataclass()
class MyDataClass2:
    my_field: str = field()
m = Mock(spec=MyDataClass2("my_field_val"))
print(m.my_field)
  • Instantiate the mock:
@dataclass()
class MyDataClass3:
    my_field: str = field()
m = Mock(spec=MyDataClass3)
m_inst = m()
print(m_inst.my_field)
  • Use a list comprehension to instantiate all fields in the spec:
@dataclass()
class MyDataClass4:
    my_field: str = field()
m = Mock(spec_set=[field.name for field in fields(MyDataClass4)])
print(m.my_field)

Upvotes: 0

Emptyless
Emptyless

Reputation: 3051

Based on the answer from mohi666. If you also want to prevent setting Mock attributes not specified in the dataclass, use spec_set instead of spec:

import mock
from dataclasses import dataclass, fields

@dataclass
class A:
   x: str


def create_dataclass_mock(obj):
    return mock.Mock(spec_set=[field.name for field in fields(obj)])

m = create_dataclass_mock(A)
m.x = 'test' # works
m.y = 'another test' # raises AttributeError

Upvotes: 2

Patrick Haugh
Patrick Haugh

Reputation: 61032

You can also pass an instance with dummy values to spec 

from unittest.mock import Mock
from dataclasses import dataclass

@dataclass
class A:
    aaa: str
    bbb: int

m = Mock(spec=A(None, None))

print(m.bbb)
# <Mock name='mock.bbb' id='139766470904856'>

Upvotes: 6

mohi666
mohi666

Reputation: 7052

I ended up using this generic helper function to achieve what spec does with regular classes:

import mock
from dataclasses import fields


def create_dataclass_mock(obj):
    return mock.Mock(spec=[field.name for field in fields(obj)])

Upvotes: 10

Related Questions