Fix unit tests

logicalclocks · Nov 1, 2024 · 4e6debb · 4e6debb
1 parent e157ad7
commit 4e6debb
Showing 1 changed file with 3 additions and 0 deletions.
diff --git a/python/tests/test_feature_group.py b/python/tests/test_feature_group.py
@@ -904,6 +904,7 @@ def test_prepare_spark_location(self, mocker, backend_fixtures):
     def test_prepare_spark_location_with_s3_connector(self, mocker, backend_fixtures):
         # Arrange
         engine = spark.Engine()
+        refetch_api = mocker.patch("hsfs.storage_connector.S3Connector.refetch")
         engine_instance = mocker.patch("hsfs.engine.get_instance", return_value=engine)
         json = backend_fixtures["feature_group"]["get_basic_info"]["response"]
         fg = feature_group.FeatureGroup.from_response_json(json)
@@ -918,13 +919,15 @@ def test_prepare_spark_location_with_s3_connector(self, mocker, backend_fixtures
         # Assert
         assert fg.location == path
         engine_instance.assert_called_once()
+        refetch_api.assert_called_once()
 
     def test_prepare_spark_location_with_s3_connector_python(
         self, mocker, backend_fixtures
     ):
         # Arrange
         engine = python.Engine()
         engine_instance = mocker.patch("hsfs.engine.get_instance", return_value=engine)
+        mocker.patch("hsfs.storage_connector.S3Connector.refetch")
         json = backend_fixtures["feature_group"]["get_basic_info"]["response"]
         fg = feature_group.FeatureGroup.from_response_json(json)
         fg._location = f"{fg.name}_{fg.version}"