Skip to content

Commit 5aa9af2

Browse files
achalsachals
authored
Teardown infrastructure after integration tests (feast-dev#1697)
* Teardown infrastructure after integration tests Signed-off-by: Achal Shah <achals@gmail.com> * Add a teardown method in feature store Signed-off-by: Achal Shah <achals@gmail.com> * fix import Signed-off-by: Achal Shah <achals@gmail.com> * Rename var to tables Signed-off-by: Achal Shah <achals@gmail.com> * Remove incorrect comment Signed-off-by: Achal Shah <achals@gmail.com>
1 parent 5957da2 commit 5aa9af2

2 files changed

Lines changed: 36 additions & 3 deletions

File tree

‎sdk/python/feast/feature_store.py‎

Lines changed: 18 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -24,6 +24,7 @@
2424
from feast import utils
2525
from feast.entity import Entity
2626
from feast.errors import FeatureNameCollisionError, FeatureViewNotFoundException
27+
from feast.feature_table import FeatureTable
2728
from feast.feature_view import FeatureView
2829
from feast.inference import (
2930
update_data_sources_with_inferred_event_timestamp_col,
@@ -254,6 +255,23 @@ def apply(
254255
partial=True,
255256
)
256257

258+
@log_exceptions_and_usage
259+
def teardown(self):
260+
tables: List[Union[FeatureView, FeatureTable]] = []
261+
feature_views = self.list_feature_views()
262+
feature_tables = self._registry.list_feature_tables(self.project)
263+
264+
tables.extend(feature_views)
265+
tables.extend(feature_tables)
266+
267+
entities = self.list_entities()
268+
269+
self._get_provider().teardown_infra(self.project, tables, entities)
270+
for feature_view in feature_views:
271+
self.delete_feature_view(feature_view.name)
272+
for feature_table in feature_tables:
273+
self._registry.delete_feature_table(feature_table.name, self.project)
274+
257275
@log_exceptions_and_usage
258276
def get_historical_features(
259277
self,

‎sdk/python/tests/test_offline_online_store_consistency.py‎

Lines changed: 18 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -111,6 +111,8 @@ def prep_bq_fs_and_fv(
111111

112112
yield fs, fv
113113

114+
fs.teardown()
115+
114116

115117
@contextlib.contextmanager
116118
def prep_local_fs_and_fv() -> Iterator[Tuple[FeatureStore, FeatureView]]:
@@ -133,10 +135,13 @@ def prep_local_fs_and_fv() -> Iterator[Tuple[FeatureStore, FeatureView]]:
133135
join_key="driver_id",
134136
value_type=ValueType.INT32,
135137
)
138+
project = f"test_local_correctness_{str(uuid.uuid4()).replace('-', '')}"
139+
print(f"Using project: {project}")
140+
136141
with tempfile.TemporaryDirectory() as repo_dir_name, tempfile.TemporaryDirectory() as data_dir_name:
137142
config = RepoConfig(
138143
registry=str(Path(repo_dir_name) / "registry.db"),
139-
project=f"test_bq_correctness_{str(uuid.uuid4()).replace('-', '')}",
144+
project=project,
140145
provider="local",
141146
online_store=SqliteOnlineStoreConfig(
142147
path=str(Path(data_dir_name) / "online_store.db")
@@ -147,6 +152,8 @@ def prep_local_fs_and_fv() -> Iterator[Tuple[FeatureStore, FeatureView]]:
147152

148153
yield fs, fv
149154

155+
fs.teardown()
156+
150157

151158
@contextlib.contextmanager
152159
def prep_redis_fs_and_fv() -> Iterator[Tuple[FeatureStore, FeatureView]]:
@@ -169,10 +176,12 @@ def prep_redis_fs_and_fv() -> Iterator[Tuple[FeatureStore, FeatureView]]:
169176
join_key="driver_id",
170177
value_type=ValueType.INT32,
171178
)
179+
project = f"test_redis_correctness_{str(uuid.uuid4()).replace('-', '')}"
180+
print(f"Using project: {project}")
172181
with tempfile.TemporaryDirectory() as repo_dir_name:
173182
config = RepoConfig(
174183
registry=str(Path(repo_dir_name) / "registry.db"),
175-
project=f"test_bq_correctness_{str(uuid.uuid4()).replace('-', '')}",
184+
project=project,
176185
provider="local",
177186
online_store=RedisOnlineStoreConfig(
178187
type="redis",
@@ -185,6 +194,8 @@ def prep_redis_fs_and_fv() -> Iterator[Tuple[FeatureStore, FeatureView]]:
185194

186195
yield fs, fv
187196

197+
fs.teardown()
198+
188199

189200
@contextlib.contextmanager
190201
def prep_dynamodb_fs_and_fv() -> Iterator[Tuple[FeatureStore, FeatureView]]:
@@ -207,10 +218,12 @@ def prep_dynamodb_fs_and_fv() -> Iterator[Tuple[FeatureStore, FeatureView]]:
207218
join_key="driver_id",
208219
value_type=ValueType.INT32,
209220
)
221+
project = f"test_dynamo_correctness_{str(uuid.uuid4()).replace('-', '')}"
222+
print(f"Using project {project}")
210223
with tempfile.TemporaryDirectory() as repo_dir_name:
211224
config = RepoConfig(
212225
registry=str(Path(repo_dir_name) / "registry.db"),
213-
project=f"test_bq_correctness_{str(uuid.uuid4()).replace('-', '')}",
226+
project=project,
214227
provider="aws",
215228
online_store=DynamoDBOnlineStoreConfig(region="us-west-2"),
216229
offline_store=FileOfflineStoreConfig(),
@@ -220,6 +233,8 @@ def prep_dynamodb_fs_and_fv() -> Iterator[Tuple[FeatureStore, FeatureView]]:
220233

221234
yield fs, fv
222235

236+
fs.teardown()
237+
223238

224239
# Checks that both offline & online store values are as expected
225240
def check_offline_and_online_features(

0 commit comments

Comments
 (0)