Correct cleanup after usage e2e tests (#2015)

pyalex · web-flow · commit 06b1f377f96e · 2021-11-09T21:43:40.000-08:00
* Correct cleanup after usage e2e tests

Signed-off-by: pyalex &lt;moskalenko.alexey@gmail.com&gt;

* do not override is_enabled directly

Signed-off-by: pyalex &lt;moskalenko.alexey@gmail.com&gt;

* reload modules after test

Signed-off-by: pyalex &lt;moskalenko.alexey@gmail.com&gt;
diff --git 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
