diff --git a/sdk/python/tests/integration/e2e/test_usage_e2e.py b/sdk/python/tests/integration/e2e/test_usage_e2e.py index c05ce8daf4e..f55fbce55cf 100644 --- a/sdk/python/tests/integration/e2e/test_usage_e2e.py +++ b/sdk/python/tests/integration/e2e/test_usage_e2e.py @@ -19,7 +19,7 @@ import pytest -from feast import Entity, RepoConfig, ValueType, usage +from feast import Entity, RepoConfig, ValueType from feast.infra.online_stores.sqlite import SqliteOnlineStoreConfig @@ -31,10 +31,18 @@ def dummy_exporter(): yield event_log -@pytest.mark.integration -def test_usage_on(dummy_exporter): - usage._is_enabled = True +@pytest.fixture(scope="function") +def enabling_toggle(): + with patch("feast.usage._is_enabled") as p: + p.__bool__.return_value = True + yield p + # return to initial state + _reload_feast() + + +@pytest.mark.integration +def test_usage_on(dummy_exporter, enabling_toggle): _reload_feast() from feast.feature_store import FeatureStore @@ -65,8 +73,8 @@ def test_usage_on(dummy_exporter): @pytest.mark.integration -def test_usage_off(dummy_exporter): - usage._is_enabled = False +def test_usage_off(dummy_exporter, enabling_toggle): + enabling_toggle.__bool__.return_value = False _reload_feast() from feast.feature_store import FeatureStore @@ -94,9 +102,7 @@ def test_usage_off(dummy_exporter): @pytest.mark.integration -def test_exception_usage_on(dummy_exporter): - usage._is_enabled = True - +def test_exception_usage_on(dummy_exporter, enabling_toggle): _reload_feast() from feast.feature_store import FeatureStore @@ -111,8 +117,8 @@ def test_exception_usage_on(dummy_exporter): @pytest.mark.integration -def test_exception_usage_off(dummy_exporter): - usage._is_enabled = False +def test_exception_usage_off(dummy_exporter, enabling_toggle): + enabling_toggle.__bool__.return_value = False _reload_feast() from feast.feature_store import FeatureStore