Use request.addfinalizer instead of the yield based approach in integ tests

sdk/python/tests/conftest.py

@@ -11,6 +11,7 @@
 # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 # See the License for the specific language governing permissions and
 # limitations under the License.
+import logging
 import multiprocessing
 from datetime import datetime, timedelta
 from sys import platform
@@ -34,6 +35,8 @@
     construct_universal_entities,
 )
 
+logger = logging.getLogger(__name__)
+
 
 def pytest_configure(config):
     if platform in ["darwin", "windows"]:
@@ -137,21 +140,33 @@ def simple_dataset_2() -> pd.DataFrame:
 @pytest.fixture(
     params=FULL_REPO_CONFIGS, scope="session", ids=[str(c) for c in FULL_REPO_CONFIGS]
 )
-def environment(request):
-    with construct_test_environment(request.param) as e:
-        yield e
+def environment(request, tmp_path_factory, worker_id):
+    logger.info("Worker: %s, Request: %s", worker_id, request.param)
+    e = construct_test_environment(request.param)
+
+    def cleanup():
+        logger.info("Running cleanup in %s, Request: %s", worker_id, request.param)
+        e.feature_store.teardown()
+
+    request.addfinalizer(cleanup)
+    return e
 
 
 @pytest.fixture()
-def local_redis_environment():
-    with construct_test_environment(
-        IntegrationTestRepoConfig(online_store=REDIS_CONFIG)
-    ) as e:
-        yield e
+def local_redis_environment(request, worker_id):
+
+    e = construct_test_environment(IntegrationTestRepoConfig(online_store=REDIS_CONFIG))
+
+    def cleanup():
+        logger.info("Running cleanup in %s, Request: %s", worker_id, request.param)
+        e.feature_store.teardown()
+
+    request.addfinalizer(cleanup)
+    return e
 
 
 @pytest.fixture(scope="session")
-def universal_data_sources(environment):
+def universal_data_sources(request, environment, tmp_path_factory, worker_id):
     entities = construct_universal_entities()
     datasets = construct_universal_datasets(
         entities, environment.start_date, environment.end_date
@@ -160,18 +175,24 @@ def universal_data_sources(environment):
         datasets, environment.data_source_creator
     )
 
-    yield entities, datasets, datasources
+    def cleanup():
+        environment.data_source_creator.teardown()
 
-    environment.data_source_creator.teardown()
+    request.addfinalizer(cleanup)
+
+    return entities, datasets, datasources
 
 
 @pytest.fixture(scope="session")
-def e2e_data_sources(environment: Environment):
+def e2e_data_sources(environment: Environment, request):
     df = create_dataset()
     data_source = environment.data_source_creator.create_data_source(
         df, environment.feature_store.project, field_mapping={"ts_1": "ts"},
     )
 
-    yield df, data_source
+    def cleanup():
+        environment.data_source_creator.teardown()
+
+    request.addfinalizer(cleanup)
 
-    environment.data_source_creator.teardown()
+    return df, data_source

sdk/python/tests/integration/feature_repos/repo_configuration.py

@@ -3,7 +3,6 @@
 import os
 import tempfile
 import uuid
-from contextlib import contextmanager
 from dataclasses import dataclass, field
 from datetime import datetime, timedelta
 from pathlib import Path
@@ -235,7 +234,6 @@ def table_name_from_data_source(ds: DataSource) -> Optional[str]:
     return None
 
 
-@contextmanager
 def construct_test_environment(
     test_repo_config: IntegrationTestRepoConfig,
     test_suite_name: str = "integration_test",
@@ -254,49 +252,47 @@ def construct_test_environment(
     offline_store_config = offline_creator.create_offline_store_config()
     online_store = test_repo_config.online_store
 
-    with tempfile.TemporaryDirectory() as repo_dir_name:
-        if test_repo_config.python_feature_server:
-            from feast.infra.feature_servers.aws_lambda.config import (
-                AwsLambdaFeatureServerConfig,
-            )
-
-            feature_server = AwsLambdaFeatureServerConfig(
-                enabled=True,
-                execution_role_name="arn:aws:iam::402087665549:role/lambda_execution_role",
-            )
-
-            registry = f"s3://feast-integration-tests/registries/{project}/registry.db"
-        else:
-            feature_server = None
-            registry = str(Path(repo_dir_name) / "registry.db")
-
-        config = RepoConfig(
-            registry=registry,
-            project=project,
-            provider=test_repo_config.provider,
-            offline_store=offline_store_config,
-            online_store=online_store,
-            repo_path=repo_dir_name,
-            feature_server=feature_server,
+    repo_dir_name = tempfile.mkdtemp()
+
+    if test_repo_config.python_feature_server:
+        from feast.infra.feature_servers.aws_lambda.config import (
+            AwsLambdaFeatureServerConfig,
         )
 
-        # Create feature_store.yaml out of the config
-        with open(Path(repo_dir_name) / "feature_store.yaml", "w") as f:
-            yaml.safe_dump(json.loads(config.json()), f)
-
-        fs = FeatureStore(repo_dir_name)
-        # We need to initialize the registry, because if nothing is applied in the test before tearing down
-        # the feature store, that will cause the teardown method to blow up.
-        fs.registry._initialize_registry()
-        environment = Environment(
-            name=project,
-            test_repo_config=test_repo_config,
-            feature_store=fs,
-            data_source_creator=offline_creator,
-            python_feature_server=test_repo_config.python_feature_server,
+        feature_server = AwsLambdaFeatureServerConfig(
+            enabled=True,
+            execution_role_name="arn:aws:iam::402087665549:role/lambda_execution_role",
         )
 
-        try:
-            yield environment
-        finally:
-            fs.teardown()
+        registry = f"s3://feast-integration-tests/registries/{project}/registry.db"
+    else:
+        feature_server = None
+        registry = str(Path(repo_dir_name) / "registry.db")
+
+    config = RepoConfig(
+        registry=registry,
+        project=project,
+        provider=test_repo_config.provider,
+        offline_store=offline_store_config,
+        online_store=online_store,
+        repo_path=repo_dir_name,
+        feature_server=feature_server,
+    )
+
+    # Create feature_store.yaml out of the config
+    with open(Path(repo_dir_name) / "feature_store.yaml", "w") as f:
+        yaml.safe_dump(json.loads(config.json()), f)
+
+    fs = FeatureStore(repo_dir_name)
+    # We need to initialize the registry, because if nothing is applied in the test before tearing down
+    # the feature store, that will cause the teardown method to blow up.
+    fs.registry._initialize_registry()
+    environment = Environment(
+        name=project,
+        test_repo_config=test_repo_config,
+        feature_store=fs,
+        data_source_creator=offline_creator,
+        python_feature_server=test_repo_config.python_feature_server,
+    )
+
+    return environment