https://prefect.io logo
Join Slack
Powered by
How to disable Prefect logger for tests?
# ask-community
d
In prefect 2, how can I test a flow or task which is using the
Copy code
get_run_logger()
which is set from the context. You can find sample test code on the thread The test keep failing with the erorr
Copy code
prefect.exceptions.MissingContextError: There is no active flow or task run context.
Copy code
import pytest
from prefect import flow, task, get_run_logger
from prefect.testing.utilities import prefect_test_harness


@pytest.fixture(autouse=True, scope="session")
def prefect_test_fixture():
    with prefect_test_harness():
        yield


@task
def sum(a: int, b: int) -> int:
    return a + b


@task
def sum_with_logger(a: int, b: int) -> int:
    logger = get_run_logger()
    logger.debug(f"sum {a}+{b}={a+b}")
    return a + b


@flow
def sum_flow_with_logger(a: int, b: int) -> int:
    return sum_with_logger(a, b)


def test_sum():
    # Test pass
    a = 1
    b = 3
    result = sum.fn(a, b)
    assert 4 == result


def test_sum_with_logger():
    # test error with prefect.exceptions.MissingContextError: There is no active flow or task run context.
    a = 1
    b = 3
    result = sum_with_logger.fn(a, b)
    assert 4 == result


def test_sum_flow_with_logger():
    # test error with prefect.exceptions.MissingContextError: There is no active flow or task run context.
    a = 1
    b = 3
    result = sum_flow_with_logger.fn(a, b)
    assert 4 == result
p
I don't actually know the pattern that we encourage for this, but I would personally just mock and monkeypatch. If we have a file 
ex.py
which has our example flow
Copy code
from prefect import flow, get_run_logger

@flow
def my_flow():
    get_run_logger().info("I am logging stuff")
    return 42
and we have our test file 
test_ex.py
I would probably do something similar to the following:
Copy code
from unittest.mock import Mock
from ex import my_flow

def test_my_flow(monkeypatch):
    mock_logger = Mock()
    monkeypatch.setattr("ex.get_run_logger", mock_logger)
    res = my_flow()
    ...
gratitude thank you 1
You could also accomplish this with a fixture
Copy code
from unittest.mock import Mock
from ex import my_flow
import pytest

@pytest.fixture
def mock_logger(monkeypatch):
    mock_logger = Mock()
    monkeypatch.setattr("ex.get_run_logger", mock_logger)
    return mock_logger

def test_my_flow(mock_logger):
    res = my_flow()
d
Try indeed... Good point👍 about the mocking in prefect 1 they use to have an helper that unit the context.
What other pattern would you use?
p
Hmmm - give me just a moment
k
Here’s the workaround that we use in `conftest.py`:
Copy code
from prefect.logging import disable_run_logger
from prefect.testing.utilities import prefect_test_harness


@pytest.fixture(autouse=True, scope="session")
def prefect_disable_logging():
    with disable_run_logger():
        yield


@pytest.fixture(autouse=True, scope="session")
def prefect_test_fixture():
    with prefect_test_harness():
        yield
gratitude thank you 2
🔥 2
This will disable the logger in your tests - not useful if you want to test the logging itself, but helpful if you don’t need to
🔥 1
p
Yeah, if you aren't worried about the output, that's a great way to go too
d
thank you @Krishnan Chandra This is working perfectly
👍 1
as I'm interesting about the logic of the task that I want to test
Thank you @Peyton Runyan for pointing me to the mock. Completly discarded this option
👍 1
Have a good Thanksgiving guys if you are in US
❤️ 3
should
Copy code
disable_run_logger
not be in
Copy code
prefect.testing.utilities
be a better place? as I don't know when you would use it outside of testing
a
true, it could be, but you can use it not only in unit tests - e.g. I can imagine someone may want to disable it for a local run (not a unit test)
92 Views
Open in Slack
PreviousNext
Powered by