diff --git a/sdk/python/tests/conftest.py b/sdk/python/tests/conftest.py index b7731fbb4c..61e591f237 100644 --- a/sdk/python/tests/conftest.py +++ b/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 diff --git a/sdk/python/tests/integration/feature_repos/repo_configuration.py b/sdk/python/tests/integration/feature_repos/repo_configuration.py index 97e379049a..6dedfb63b2 100644 --- a/sdk/python/tests/integration/feature_repos/repo_configuration.py +++ b/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 diff --git a/sdk/python/tests/integration/registration/test_universal_types.py b/sdk/python/tests/integration/registration/test_universal_types.py index 7ef1c5c1e7..3716d8f870 100644 --- a/sdk/python/tests/integration/registration/test_universal_types.py +++ b/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