How do I patch a method registered by a decorator in Python's datashape?

Question

I'm using the datashape Python package and registering a new type with the @datashape.discover.register decorator. I'd like to test that when I call datashape.discover on an object of the type I'm registering, it calls the function being decorated. I'd also like to do this with good unit testing principles, meaning not actually executing the function being decorated, as it would have side effects I don't want in the test. However, this isn't working.

Here's some sample code to demonstrate the problem:

myfile.py:

@datashape.discover.register(SomeType)
def discover_some_type(data)
    ...some stuff i don't want done in a unit test...

test_myfile.py:

class TestDiscoverSomeType(unittest.TestCase):
    @patch('myfile.discover_some_type')
    def test_discover_some_type(self, mock_discover_some_type):
        file_to_discover = SomeType()

        datashape.discover(file_to_discover)

        mock_discover_some_type.assert_called_with(file_to_discover)

The issue seems to be that the function I want mocked is mocked in the body of the test, however, it was not mocked when it was decorated (i.e. when it was imported). The discover.register function essentially internally registers the function being decorated to look it up when discover() is called with an argument of the given type. Unfortunately, it seems to internally register the real function every time, and not the patched version I want, so it will always call the real function.

Any thoughts on how to be able to patch the function being decorated and assert that it is called when datashape.discover is called?

Answer 1

Here's a solution I've found that's only a little hacky:

sometype.py:

def discover_some_type(data):
    ...some stuff i don't want done in a unit test...

discovery_channel.py:

import sometype

@datashape.discover.register(SomeType)
def discover_some_type(data):
    return sometype.discover_some_type(data)

test_sometype.py:

class TestDiscoverSomeType(unittest.TestCase):
    @patch('sometype.discover_some_type')
    def test_discover_some_type(self, mock_discover_some_type):
        import discovery_channel

        file_to_discover = SomeType()

        datashape.discover(file_to_discover)

        mock_discover_some_type.assert_called_with(file_to_discover)

The key is that you have to patch out whatever will actually do stuff before you import the module that has the decorated function that will register the patched function to datashape. This unfortunately means that you can't have your decorated function and the function doing the discovery in the same module (so things that should logically go together are now apart). And you have the somewhat hacky import-in-a-function in your unit test (to trigger the discover.register). But at least it works.