CODEWITHSUNDEEP
pythonunit-testingmockingpython-unittestpython-mock
Joshmaker
Joshmaker

Reputation: 4248

How to combine inheritance with mock autospec in Python

The Problem

I'm looking for a way to properly Mock objects in my unit tests, but I'm having trouble getting unittest.mock.create_autospec or unittest.mock.Mock to do what I need. I think I need to use inheritance with a Mock object, but I'm having trouble making it work.

Image that I need to mock the classes in a 3rd party module that looks something like this (pretend the lines with raise NotImplementedError are external API calls I want to avoid in my unit tests):

class FileStorageBucket():
    def __init__(self, bucketname: str) -> None:
        self.bucketname = bucketname

    def download(self, filename) -> None:
        raise NotImplementedError

    # ...lots more methods...


class FileStorageClient():
    def auth(self, username: str, password: str) -> None:
        raise NotImplementedError

    def get_bucket(self, bucketname: str) -> FileStorageBucket:
        raise NotImplementedError
        return FileStorageBucket(bucketname)

    # ...lots more methods...

And it might be used elsewhere in my application like this:

client = FileStorageClient()
client.auth("me", "mypassword")
client.get_bucket("my-bucket").download("my-file.jpg")

If I replace FileStorageClient with a Mock object, I'd like to be able to find out if my unit tests run any code where:

So, client.get_bucket("foo").download() should raise an exception that filename is a required argument for .download().

Things I've tried:

First, I tried using create_autospec. It is able to catch some types of errors:

>>> MockClient = create_autospec(FileStorageClient)
>>> client = MockClient()
>>> client.auth(user_name="name", password="password")
TypeError: missing a required argument: 'username'

But, of course, because it doesn't know the return type that get_bucket should have it doesn't catch other types of errors:

>>> MockClient = create_autospec(FileStorageClient)
>>> client = MockClient()
>>> client.get_bucket("foo").download(wrong_arg="foo")
<MagicMock name='mock.get_bucket().download()' id='4554265424'>

I thought I could solve this by creating classes that inherited from the create_autospec output:

class MockStorageBucket(create_autospec(FileStorageBucket)):
    def path(self, filename) -> str:
        return f"/{self.bucketname}/{filename}"


class MockStorageClient(create_autospec(FileStorageClient)):
    def get_bucket(self, bucketname: str):
        bucket = MockStorageBucket()
        bucket.bucketname = bucketname
        return bucket

But it doesn't actually return a MockStorageBucket instance as expected:

>>> client = MockStorageClient()
>>> client.get_bucket("foo").download(wrong_arg="foo")
<MagicMock name='mock.get_bucket().download()' id='4554265424'>

So then I tried inheriting from Mock and manually setting the "spec" in the init:

class MockStorageBucket(Mock):
    def __init__(self, *args, **kwargs):
        # Pass `FileStorageBucket` as the "spec"
        super().__init__(FileStorageBucket, *args, **kwargs)

    def path(self, filename) -> str:
        return f"/{self.bucketname}/{filename}"


class MockStorageClient(Mock):
    def __init__(self, *args, **kwargs):
        # Pass `FileStorageClient` as the "spec"
        super().__init__(FileStorageClient, *args, **kwargs)

    def get_bucket(self, bucketname: str):
        bucket = MockStorageBucket()
        bucket.bucketname = bucketname
        return bucket

Now, the get_bucket method returns a MockStorageBucket instance as expected, and I am able to catch some errors, such as accessing attributes that don't exist:

>>> client = MockStorageClient()
>>> client.get_bucket("my-bucket")
<__main__.FileStorageBucket at 0x10f7a0110>
>>> client.get_bucket("my-bucket").foobar
AttributeError: Mock object has no attribute 'foobar'

However, unlike the Mock instance created with create_autospec Mock instances with Mock(spec=whatever) don't appear to check that the correct arguments are passed to a function:

>>> client.auth(wrong_arg=1)
<__main__.FileStorageClient at 0x10dac5990>

Upvotes: 1

Views: 1679

Answers (2)

Joshmaker
Joshmaker

Reputation: 4248

I think the full code I want is something like:

def mock_client_factory() -> Mock:
    MockClient = create_autospec(FileStorageClient)

    def mock_bucket_factory(bucketname: str) -> Mock:
        MockBucket = create_autospec(FileStorageBucket)
        mock_bucket = MockBucket(bucketname=bucketname)
        mock_bucket.bucketname = bucketname
        return mock_bucket

    mock_client = MockClient()
    mock_client.get_bucket.side_effect = mock_bucket_factory
    return mock_client

Upvotes: 0

nathanielobrown
nathanielobrown

Reputation: 1012

Just set the return_value on your get_bucket method to be another mock with a different spec. You don't need to mess around with creating MockStorageBucket and MockStorageClient.

mock_client = create_autospec(FileStorageClient, spec_set=True)
mock_bucket = create_autospec(FileStorageBucket, spec_set=True)
mock_client.get_bucket.return_value = mock_bucket

mock_client.get_bucket("my-bucket").download("my-file.jpg")

Upvotes: 3

Related Questions