https://prefect.io logo
j

Joël Luijmes

02/01/2021, 3:27 PM
I’m using a Context Manager, but how can I mock the behavior of Prefect such that I can only test the setup function? I tried
Copy code
@pytest.fixture
def mock_resource_manager(monkeypatch):
    mock = MagicMock(return_value=None)
    monkeypatch.setattr("prefect.tasks.core.resource_manager.ResourceManager.__call__", mock)
    return mock
But this gives weird errors down the road, as I also want to mock prefect.config/context for other purposes.
1
j

Josh Greenhalgh

02/01/2021, 3:30 PM
Don't? Wrap your actual logic in a function that pulls from the context and pass into an interior function that is pure and only depends in its args?
Then unit test the interior function?
From a pure pytest perspective have you tried using
setattr
on the actual import;
Copy code
# flow.py

from prefect.tasks.core.resource_manager import ResourceManager

...

# test_flow.py
@pytest.fixture
def mock_resource_manager(monkeypatch):
    mock = MagicMock(return_value=None)
    monkeypatch.setattr("flow.ResourceManager.__call__", mock)
    return mock
j

Joël Luijmes

02/02/2021, 2:02 PM
Hmm not sure if I understtand you. Let me give some more details, I got this resoure manager of which I want to test the setup function
Copy code
@resource_manager
class CloudSQLProxyManager:
    def __init__(self, name, instance, keyfile):
        self.instance = instance
        self.keyfile = keyfile
        self.name = name
        self.skip_deployment = prefect.config.get('environment', 'development').lower() != 'production'

    def setup(self):
        if self.skip_deployment:
            <http://prefect.context.logger.info|prefect.context.logger.info>(f"Skipping creation of CloudSQL Proxy")
            raise signals.SUCCESS("Creation of CloudSQL Proxy skipped")
        # ....
I tried your suggestion, without mocking:
Copy code
`def test__setup__skips_on_non_production(mock_resource_manager):
    prefect.config['environment'] = 'development'

    flow = Flow("unit test")
    manager = CloudSQLProxyManager("name", "instance", "keyfile", flow=flow)
    with pytest.raises(signals.SUCCESS, match="Creation of CloudSQL Proxy skipped"):
        print(manager)
        manager.setup(manager)
But then manager is an instance of ResourceManager (and not CloudSQLLProxyManager) I tried mocking the call (as shown in my initial post), but that makes manager None.
Okay found a solution! What I wanted is to have the original class, without the resource_manager decorator applied. I managed to achieve this by creating the following fixture
Copy code
@pytest.fixture
def CloudSQLProxyManager(monkeypatch):
    monkeypatch.setattr("prefect.tasks.core.resource_manager.ResourceManager.__call__", lambda x: x)
    
    from src.modules.cloudsql_proxy_manager import CloudSQLProxyManager
    return CloudSQLProxyManager().resource_class
This fixture returns the original class - without decorations - and now I can finally unit test my setup function 🙂
🎉 2
2 Views