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

… tests (#2089)

* Use request.addfinalizer instead of the yield based approach

Signed-off-by: Achal Shah <achals@gmail.com>

* remove unused fixtures

Signed-off-by: Achal Shah <achals@gmail.com>

* Remove yields from the type tests

Signed-off-by: Achal Shah <achals@gmail.com>

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"]:
@@ -138,20 +141,29 @@ def simple_dataset_2() -> pd.DataFrame:
     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
+    e = construct_test_environment(request.param)
+
+    def cleanup():
+        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():
+        e.feature_store.teardown()
+
+    request.addfinalizer(cleanup)
+    return e
 
 
 @pytest.fixture(scope="session")
-def universal_data_sources(environment):
+def universal_data_sources(request, environment):
     entities = construct_universal_entities()
     datasets = construct_universal_datasets(
         entities, environment.start_date, environment.end_date
@@ -160,18 +172,24 @@ def universal_data_sources(environment):
         datasets, environment.data_source_creator
     )
 
-    yield entities, datasets, datasources
+    def cleanup():
+        # logger.info("Running cleanup in %s, Request: %s", worker_id, request.param)
+        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

sdk/python/tests/integration/registration/test_universal_types.py

@@ -71,7 +71,7 @@ class TypeTestConfig:
     ids=[str(c) for c in OFFLINE_TYPE_TEST_CONFIGS],
 )
 def offline_types_test_fixtures(request):
-    yield from get_fixtures(request)
+    return get_fixtures(request)
 
 
 @pytest.fixture(
@@ -80,7 +80,7 @@ def offline_types_test_fixtures(request):
     ids=[str(c) for c in ONLINE_TYPE_TEST_CONFIGS],
 )
 def online_types_test_fixtures(request):
-    yield from get_fixtures(request)
+    return get_fixtures(request)
 
 
 def get_fixtures(request):
@@ -89,30 +89,36 @@ def get_fixtures(request):
     test_project_id = f"{config.entity_type}{config.feature_dtype}{config.feature_is_list}".replace(
         ".", ""
     ).lower()
-    with construct_test_environment(
+    type_test_environment = construct_test_environment(
         test_repo_config=config.test_repo_config,
         test_suite_name=f"test_{test_project_id}",
-    ) as type_test_environment:
-        config = request.param
-        df = create_dataset(
-            config.entity_type,
-            config.feature_dtype,
-            config.feature_is_list,
-            config.has_empty_list,
-        )
-        data_source = type_test_environment.data_source_creator.create_data_source(
-            df,
-            destination_name=type_test_environment.feature_store.project,
-            field_mapping={"ts_1": "ts"},
-        )
-        fv = create_feature_view(
-            config.feature_dtype,
-            config.feature_is_list,
-            config.has_empty_list,
-            data_source,
-        )
-        yield type_test_environment, config, data_source, fv
+    )
+    config = request.param
+    df = create_dataset(
+        config.entity_type,
+        config.feature_dtype,
+        config.feature_is_list,
+        config.has_empty_list,
+    )
+    data_source = type_test_environment.data_source_creator.create_data_source(
+        df,
+        destination_name=type_test_environment.feature_store.project,
+        field_mapping={"ts_1": "ts"},
+    )
+    fv = create_feature_view(
+        config.feature_dtype,
+        config.feature_is_list,
+        config.has_empty_list,
+        data_source,
+    )
+
+    def cleanup():
         type_test_environment.data_source_creator.teardown()
+        type_test_environment.feature_store.teardown()
+
+    request.addfinalizer(cleanup)
+
+    return type_test_environment, config, data_source, fv
 
 
 @pytest.mark.integration