updated tests

Signed-off-by: Francisco Javier Arceo <[email protected]>

sdk/python/tests/unit/online_store/test_online_writes.py

@@ -77,6 +77,7 @@ def setUp(self):
             )
             # Before apply() join_keys is empty
             assert driver_stats_fv.join_keys == []
+            assert driver_stats_fv.entity_columns == []
 
             @on_demand_feature_view(
                 sources=[driver_stats_fv[["conv_rate", "acc_rate"]]],
@@ -104,6 +105,8 @@ def test_view(inputs: dict[str, Any]) -> dict[str, Any]:
             )
             # after apply() join_keys is [driver]
             assert driver_stats_fv.join_keys == [driver.join_key]
+            assert driver_stats_fv.entity_columns[0].name == driver.join_key
+
             self.store.write_to_online_store(
                 feature_view_name="driver_hourly_stats", df=driver_df
             )

sdk/python/tests/unit/test_on_demand_python_transformation.py

@@ -615,7 +615,7 @@ def python_stored_writes_feature_view(
 
             # Note here that after apply() is called, the entity_columns are populated with the join_key
             # assert fv_applied.entity_columns[0].name == driver.join_key
-            assert fv_applied.entity_columns == []
+            assert fv_applied.entity_columns[0].name == driver.join_key
             assert odfv_applied.entity_columns[0].name == driver.join_key
 
             assert len(self.store.list_all_feature_views()) == 2