Teardown infrastructure after integration tests (#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>
feast-dev · Jul 13, 2021 · 5aa9af2 · 5aa9af2
1 parent 5957da2
commit 5aa9af2
Show file tree

Hide file tree

Showing 2 changed files with 36 additions and 3 deletions.
diff --git a/sdk/python/feast/feature_store.py b/sdk/python/feast/feature_store.py
@@ -24,6 +24,7 @@
 from feast import utils
 from feast.entity import Entity
 from feast.errors import FeatureNameCollisionError, FeatureViewNotFoundException
+from feast.feature_table import FeatureTable
 from feast.feature_view import FeatureView
 from feast.inference import (
     update_data_sources_with_inferred_event_timestamp_col,
@@ -254,6 +255,23 @@ def apply(
             partial=True,
         )
 
+    @log_exceptions_and_usage
+    def teardown(self):
+        tables: List[Union[FeatureView, FeatureTable]] = []
+        feature_views = self.list_feature_views()
+        feature_tables = self._registry.list_feature_tables(self.project)
+
+        tables.extend(feature_views)
+        tables.extend(feature_tables)
+
+        entities = self.list_entities()
+
+        self._get_provider().teardown_infra(self.project, tables, entities)
+        for feature_view in feature_views:
+            self.delete_feature_view(feature_view.name)
+        for feature_table in feature_tables:
+            self._registry.delete_feature_table(feature_table.name, self.project)
+
     @log_exceptions_and_usage
     def get_historical_features(
         self,

diff --git a/sdk/python/tests/test_offline_online_store_consistency.py b/sdk/python/tests/test_offline_online_store_consistency.py
@@ -111,6 +111,8 @@ def prep_bq_fs_and_fv(
 
         yield fs, fv
 
+        fs.teardown()
+
 
 @contextlib.contextmanager
 def prep_local_fs_and_fv() -> Iterator[Tuple[FeatureStore, FeatureView]]:
@@ -133,10 +135,13 @@ def prep_local_fs_and_fv() -> Iterator[Tuple[FeatureStore, FeatureView]]:
             join_key="driver_id",
             value_type=ValueType.INT32,
         )
+        project = f"test_local_correctness_{str(uuid.uuid4()).replace('-', '')}"
+        print(f"Using project: {project}")
+
         with tempfile.TemporaryDirectory() as repo_dir_name, tempfile.TemporaryDirectory() as data_dir_name:
             config = RepoConfig(
                 registry=str(Path(repo_dir_name) / "registry.db"),
-                project=f"test_bq_correctness_{str(uuid.uuid4()).replace('-', '')}",
+                project=project,
                 provider="local",
                 online_store=SqliteOnlineStoreConfig(
                     path=str(Path(data_dir_name) / "online_store.db")
@@ -147,6 +152,8 @@ def prep_local_fs_and_fv() -> Iterator[Tuple[FeatureStore, FeatureView]]:
 
             yield fs, fv
 
+            fs.teardown()
+
 
 @contextlib.contextmanager
 def prep_redis_fs_and_fv() -> Iterator[Tuple[FeatureStore, FeatureView]]:
@@ -169,10 +176,12 @@ def prep_redis_fs_and_fv() -> Iterator[Tuple[FeatureStore, FeatureView]]:
             join_key="driver_id",
             value_type=ValueType.INT32,
         )
+        project = f"test_redis_correctness_{str(uuid.uuid4()).replace('-', '')}"
+        print(f"Using project: {project}")
         with tempfile.TemporaryDirectory() as repo_dir_name:
             config = RepoConfig(
                 registry=str(Path(repo_dir_name) / "registry.db"),
-                project=f"test_bq_correctness_{str(uuid.uuid4()).replace('-', '')}",
+                project=project,
                 provider="local",
                 online_store=RedisOnlineStoreConfig(
                     type="redis",
@@ -185,6 +194,8 @@ def prep_redis_fs_and_fv() -> Iterator[Tuple[FeatureStore, FeatureView]]:
 
             yield fs, fv
 
+            fs.teardown()
+
 
 @contextlib.contextmanager
 def prep_dynamodb_fs_and_fv() -> Iterator[Tuple[FeatureStore, FeatureView]]:
@@ -207,10 +218,12 @@ def prep_dynamodb_fs_and_fv() -> Iterator[Tuple[FeatureStore, FeatureView]]:
             join_key="driver_id",
             value_type=ValueType.INT32,
         )
+        project = f"test_dynamo_correctness_{str(uuid.uuid4()).replace('-', '')}"
+        print(f"Using project {project}")
         with tempfile.TemporaryDirectory() as repo_dir_name:
             config = RepoConfig(
                 registry=str(Path(repo_dir_name) / "registry.db"),
-                project=f"test_bq_correctness_{str(uuid.uuid4()).replace('-', '')}",
+                project=project,
                 provider="aws",
                 online_store=DynamoDBOnlineStoreConfig(region="us-west-2"),
                 offline_store=FileOfflineStoreConfig(),
@@ -220,6 +233,8 @@ def prep_dynamodb_fs_and_fv() -> Iterator[Tuple[FeatureStore, FeatureView]]:
 
             yield fs, fv
 
+            fs.teardown()
+
 
 # Checks that both offline & online store values are as expected
 def check_offline_and_online_features(