|
7 | 7 | from feast import Entity, FeatureService, FeatureView, Field, FileSource
|
8 | 8 | from feast.api.registry.rest.rest_registry_server import RestRegistryServer
|
9 | 9 | from feast.feature_store import FeatureStore
|
| 10 | +from feast.infra.offline_stores.file_source import SavedDatasetFileStorage |
10 | 11 | from feast.repo_config import RepoConfig
|
| 12 | +from feast.saved_dataset import SavedDataset |
11 | 13 | from feast.types import Float64, Int64
|
12 | 14 | from feast.value_type import ValueType
|
13 | 15 |
|
@@ -67,8 +69,19 @@ def fastapi_test_app():
|
67 | 69 | features=[user_profile_feature_view],
|
68 | 70 | )
|
69 | 71 |
|
| 72 | + # Create a saved dataset for testing |
| 73 | + saved_dataset_storage = SavedDatasetFileStorage(path=parquet_file_path) |
| 74 | + test_saved_dataset = SavedDataset( |
| 75 | + name="test_saved_dataset", |
| 76 | + features=["user_profile:age", "user_profile:income"], |
| 77 | + join_keys=["user_id"], |
| 78 | + storage=saved_dataset_storage, |
| 79 | + tags={"environment": "test", "version": "1.0"}, |
| 80 | + ) |
| 81 | + |
70 | 82 | # Apply objects
|
71 | 83 | store.apply([user_id_entity, user_profile_feature_view, user_feature_service])
|
| 84 | + store._registry.apply_saved_dataset(test_saved_dataset, "demo_project") |
72 | 85 |
|
73 | 86 | # Build REST app with registered routes
|
74 | 87 | rest_server = RestRegistryServer(store)
|
@@ -248,3 +261,69 @@ def test_lineage_endpoint_error_handling(fastapi_test_app):
|
248 | 261 | # Test object relationships with missing parameters
|
249 | 262 | response = fastapi_test_app.get("/lineage/objects/featureView/test_fv")
|
250 | 263 | assert response.status_code == 422 # Missing required project parameter
|
| 264 | + |
| 265 | + |
| 266 | +def test_saved_datasets_via_rest(fastapi_test_app): |
| 267 | + # Test list saved datasets endpoint |
| 268 | + response = fastapi_test_app.get("/saved_datasets?project=demo_project") |
| 269 | + assert response.status_code == 200 |
| 270 | + response_data = response.json() |
| 271 | + assert "saved_datasets" in response_data |
| 272 | + assert isinstance(response_data["saved_datasets"], list) |
| 273 | + assert len(response_data["saved_datasets"]) == 1 |
| 274 | + |
| 275 | + saved_dataset = response_data["saved_datasets"][0] |
| 276 | + assert saved_dataset["spec"]["name"] == "test_saved_dataset" |
| 277 | + assert "user_profile:age" in saved_dataset["spec"]["features"] |
| 278 | + assert "user_profile:income" in saved_dataset["spec"]["features"] |
| 279 | + assert "user_id" in saved_dataset["spec"]["joinKeys"] |
| 280 | + assert saved_dataset["spec"]["tags"]["environment"] == "test" |
| 281 | + assert saved_dataset["spec"]["tags"]["version"] == "1.0" |
| 282 | + |
| 283 | + # Test get specific saved dataset endpoint |
| 284 | + response = fastapi_test_app.get( |
| 285 | + "/saved_datasets/test_saved_dataset?project=demo_project" |
| 286 | + ) |
| 287 | + assert response.status_code == 200 |
| 288 | + response_data = response.json() |
| 289 | + assert response_data["spec"]["name"] == "test_saved_dataset" |
| 290 | + assert "user_profile:age" in response_data["spec"]["features"] |
| 291 | + assert "user_profile:income" in response_data["spec"]["features"] |
| 292 | + |
| 293 | + # Test with allow_cache parameter |
| 294 | + response = fastapi_test_app.get( |
| 295 | + "/saved_datasets/test_saved_dataset?project=demo_project&allow_cache=false" |
| 296 | + ) |
| 297 | + assert response.status_code == 200 |
| 298 | + assert response.json()["spec"]["name"] == "test_saved_dataset" |
| 299 | + |
| 300 | + # Test with tags filter |
| 301 | + response = fastapi_test_app.get( |
| 302 | + "/saved_datasets?project=demo_project&tags=environment:test" |
| 303 | + ) |
| 304 | + assert response.status_code == 200 |
| 305 | + assert len(response.json()["saved_datasets"]) == 1 |
| 306 | + |
| 307 | + # Test with non-matching tags filter |
| 308 | + response = fastapi_test_app.get( |
| 309 | + "/saved_datasets?project=demo_project&tags=environment:production" |
| 310 | + ) |
| 311 | + assert response.status_code == 200 |
| 312 | + assert len(response.json()["saved_datasets"]) == 0 |
| 313 | + |
| 314 | + # Test with multiple tags filter |
| 315 | + response = fastapi_test_app.get( |
| 316 | + "/saved_datasets?project=demo_project&tags=environment:test&tags=version:1.0" |
| 317 | + ) |
| 318 | + assert response.status_code == 200 |
| 319 | + assert len(response.json()["saved_datasets"]) == 1 |
| 320 | + |
| 321 | + # Test non-existent saved dataset |
| 322 | + response = fastapi_test_app.get("/saved_datasets/non_existent?project=demo_project") |
| 323 | + assert response.status_code == 404 |
| 324 | + |
| 325 | + # Test missing project parameter |
| 326 | + response = fastapi_test_app.get("/saved_datasets/test_saved_dataset") |
| 327 | + assert ( |
| 328 | + response.status_code == 422 |
| 329 | + ) # Unprocessable Entity for missing required query param |
0 commit comments