diff --git a/sdk/python/feast/ui_server.py b/sdk/python/feast/ui_server.py index 6883dc1105e..e85a11afc06 100644 --- a/sdk/python/feast/ui_server.py +++ b/sdk/python/feast/ui_server.py @@ -4,7 +4,7 @@ from typing import Callable, Optional import uvicorn -from fastapi import FastAPI, Response +from fastapi import FastAPI, Response, status from fastapi.middleware.cors import CORSMiddleware from fastapi.staticfiles import StaticFiles from pydantic import BaseModel @@ -76,12 +76,22 @@ def shutdown_event(): @app.get("/registry") def read_registry(): if registry_proto is None: - return Response(status_code=503) # Service Unavailable + return Response( + status_code=status.HTTP_503_SERVICE_UNAVAILABLE + ) # Service Unavailable return Response( content=registry_proto.SerializeToString(), media_type="application/octet-stream", ) + @app.get("/health") + def health(): + return ( + Response(status_code=status.HTTP_200_OK) + if registry_proto + else Response(status_code=status.HTTP_503_SERVICE_UNAVAILABLE) + ) + @app.post("/save-document") async def save_document_endpoint(request: SaveDocumentRequest): try: diff --git a/sdk/python/tests/unit/test_ui_server.py b/sdk/python/tests/unit/test_ui_server.py new file mode 100644 index 00000000000..c5a85a85382 --- /dev/null +++ b/sdk/python/tests/unit/test_ui_server.py @@ -0,0 +1,297 @@ +import contextlib +import json +import os +import tempfile +from pathlib import Path +from unittest.mock import MagicMock, patch + +import assertpy +import pytest +from fastapi.testclient import TestClient + +from feast.ui_server import get_app + +# Test constants +EXPECTED_SUCCESS_STATUS = 200 +EXPECTED_ERROR_STATUS = 503 +TEST_PROJECT_NAME = "test_project" +REGISTRY_TTL_SECS = 60 + + +def _create_mock_ui_files(temp_dir): + """Helper function to create required UI files structure""" + ui_dir = os.path.join(temp_dir, "ui", "build") + os.makedirs(ui_dir, exist_ok=True) + + # Create projects-list.json file + projects_file = os.path.join(ui_dir, "projects-list.json") + with open(projects_file, "w") as f: + json.dump({"projects": []}, f) + + # Create index.html file + index_file = os.path.join(ui_dir, "index.html") + with open(index_file, "w") as f: + f.write("Test UI") + + +@contextlib.contextmanager +def _setup_importlib_mocks(temp_dir): + """Helper function to setup importlib resource mocks. + + This function mocks the importlib_resources functionality used by the UI server + to serve static files. It creates a proper context manager that returns the + temporary directory path when used with importlib_resources.as_file(). + """ + mock_path = Path(temp_dir) + + # Create a proper context manager mock + mock_context_manager = MagicMock() + mock_context_manager.__enter__.return_value = mock_path + mock_context_manager.__exit__.return_value = None + + # Mock the files() method to return a mock that supports division + mock_file_ref = MagicMock() + mock_file_ref.__truediv__.return_value = MagicMock() + + with ( + patch("feast.ui_server.importlib_resources.files") as mock_files, + patch("feast.ui_server.importlib_resources.as_file") as mock_as_file, + ): + mock_files.return_value = mock_file_ref + mock_as_file.return_value = mock_context_manager + + yield mock_files, mock_as_file + + +@pytest.fixture +def mock_feature_store(): + """Fixture for creating a mock feature store""" + mock_store = MagicMock() + mock_store.refresh_registry = MagicMock() + return mock_store + + +@pytest.fixture +def ui_app_with_registry(mock_feature_store): + """Fixture for UI app with valid registry data. + + Creates a UI app instance with a properly configured feature store + that has valid registry data available for testing endpoints that + require registry access. + """ + mock_registry = MagicMock() + mock_proto = MagicMock() + mock_proto.SerializeToString.return_value = b"mock_proto_data" + mock_registry.proto.return_value = mock_proto + mock_feature_store.registry = mock_registry + + with tempfile.TemporaryDirectory() as temp_dir: + _create_mock_ui_files(temp_dir) + + with _setup_importlib_mocks(temp_dir): + app = get_app(mock_feature_store, TEST_PROJECT_NAME, REGISTRY_TTL_SECS) + yield app + + +@pytest.fixture +def ui_app_without_registry(mock_feature_store): + """Fixture for UI app with None registry data. + + Creates a UI app instance with a feature store that has no registry + data available, used for testing error conditions and service + unavailable responses. + """ + mock_registry = MagicMock() + mock_registry.proto.return_value = None + mock_feature_store.registry = mock_registry + + with tempfile.TemporaryDirectory() as temp_dir: + _create_mock_ui_files(temp_dir) + + with _setup_importlib_mocks(temp_dir): + app = get_app(mock_feature_store, TEST_PROJECT_NAME, REGISTRY_TTL_SECS) + yield app + + +def test_ui_server_health_endpoint(ui_app_with_registry): + """Test the UI server health endpoint returns 200 when registry is available. + + This test verifies that the /health endpoint correctly returns HTTP 200 + when the feature store registry is properly initialized and contains data. + """ + client = TestClient(ui_app_with_registry) + response = client.get("/health") + assertpy.assert_that(response.status_code).is_equal_to(EXPECTED_SUCCESS_STATUS) + + +def test_ui_server_health_endpoint_with_none_registry(ui_app_without_registry): + """Test the UI server health endpoint returns 503 when registry is None. + + This test verifies that the /health endpoint correctly returns HTTP 503 + (Service Unavailable) when the feature store registry is not available + or contains no data. + """ + client = TestClient(ui_app_without_registry) + response = client.get("/health") + assertpy.assert_that(response.status_code).is_equal_to(EXPECTED_ERROR_STATUS) + + +def test_registry_endpoint_with_valid_data(ui_app_with_registry): + """Test the registry endpoint returns valid data with correct content type. + + This test verifies that the /registry endpoint correctly returns HTTP 200 + with the proper content-type header when registry data is available. + """ + client = TestClient(ui_app_with_registry) + response = client.get("/registry") + assertpy.assert_that(response.status_code).is_equal_to(EXPECTED_SUCCESS_STATUS) + assertpy.assert_that(response.headers["content-type"]).is_equal_to( + "application/octet-stream" + ) + + +def test_registry_endpoint_with_none_data(ui_app_without_registry): + """Test the registry endpoint returns 503 when registry data is None. + + This test verifies that the /registry endpoint correctly returns HTTP 503 + (Service Unavailable) when no registry data is available. + """ + client = TestClient(ui_app_without_registry) + response = client.get("/registry") + assertpy.assert_that(response.status_code).is_equal_to(EXPECTED_ERROR_STATUS) + + +def test_save_document_endpoint_success(ui_app_with_registry): + """Test the save document endpoint successfully saves data to a labels file. + + This test verifies that the /save-document endpoint correctly processes + a valid request, creates a labels file, and returns success confirmation. + """ + client = TestClient(ui_app_with_registry) + + # Create a temporary file in the current working directory for testing + with tempfile.NamedTemporaryFile( + mode="w", suffix=".py", delete=False, dir=os.getcwd() + ) as f: + test_file_path = f.name + f.write("# Test file content") + + try: + request_data = { + "file_path": test_file_path, + "data": {"test": "data", "key": "value"}, + } + + response = client.post("/save-document", json=request_data) + assertpy.assert_that(response.status_code).is_equal_to(EXPECTED_SUCCESS_STATUS) + + response_data = response.json() + assertpy.assert_that(response_data["success"]).is_true() + assertpy.assert_that(response_data).contains_key("saved_to") + + # Verify the file was created + labels_file = response_data["saved_to"] + assertpy.assert_that(os.path.exists(labels_file)).is_true() + + with open(labels_file, "r") as f: + saved_data = json.load(f) + assertpy.assert_that(saved_data).is_equal_to(request_data["data"]) + + finally: + # Cleanup + if os.path.exists(test_file_path): + os.unlink(test_file_path) + labels_file = test_file_path.replace(".py", "-labels.json") + if os.path.exists(labels_file): + os.unlink(labels_file) + + +def test_save_document_endpoint_invalid_path(ui_app_with_registry): + """Test the save document endpoint returns error for invalid file path. + + This test verifies that the /save-document endpoint correctly rejects + file paths that are outside the current working directory for security. + """ + client = TestClient(ui_app_with_registry) + + request_data = { + "file_path": "/invalid/absolute/path/outside/workspace.py", + "data": {"test": "data"}, + } + + response = client.post("/save-document", json=request_data) + assertpy.assert_that(response.status_code).is_equal_to(EXPECTED_SUCCESS_STATUS) + + response_data = response.json() + assertpy.assert_that(response_data).contains_key("error") + assertpy.assert_that(response_data["error"]).contains("Invalid file path") + + +def test_save_document_endpoint_exception_handling(ui_app_with_registry): + """Test the save document endpoint handles exceptions gracefully. + + This test verifies that the /save-document endpoint properly catches + and returns error responses when exceptions occur during processing. + """ + client = TestClient(ui_app_with_registry) + + # Test with a file path outside the current working directory (will cause an exception) + request_data = { + "file_path": "/invalid/absolute/path/outside/workspace.py", + "data": {"test": "data"}, + } + + response = client.post("/save-document", json=request_data) + assertpy.assert_that(response.status_code).is_equal_to(EXPECTED_SUCCESS_STATUS) + + response_data = response.json() + assertpy.assert_that(response_data).contains_key("error") + assertpy.assert_that(response_data["error"]).contains("Invalid file path") + + +@pytest.mark.parametrize( + "registry_available,expected_status", + [(True, EXPECTED_SUCCESS_STATUS), (False, EXPECTED_ERROR_STATUS)], +) +def test_health_endpoint_status( + registry_available, expected_status, mock_feature_store +): + """Test the health endpoint returns correct status based on registry availability. + + This parametrized test verifies that the /health endpoint returns the + appropriate HTTP status code based on whether registry data is available. + """ + if registry_available: + mock_registry = MagicMock() + mock_proto = MagicMock() + mock_proto.SerializeToString.return_value = b"mock_proto_data" + mock_registry.proto.return_value = mock_proto + mock_feature_store.registry = mock_registry + else: + mock_registry = MagicMock() + mock_registry.proto.return_value = None + mock_feature_store.registry = mock_registry + + with tempfile.TemporaryDirectory() as temp_dir: + _create_mock_ui_files(temp_dir) + + with _setup_importlib_mocks(temp_dir): + app = get_app(mock_feature_store, TEST_PROJECT_NAME, REGISTRY_TTL_SECS) + client = TestClient(app) + response = client.get("/health") + assertpy.assert_that(response.status_code).is_equal_to(expected_status) + + +def test_catch_all_route(ui_app_with_registry): + """Test the catch-all route for React router paths. + + This test reveals a bug in the original UI server code where ui_dir + is not in scope for the catch_all function. The ui_dir variable is defined + inside the importlib_resources context manager but used outside of it. + This causes a NameError when the route is accessed. + """ + client = TestClient(ui_app_with_registry) + + # The route will fail due to the scope issue with ui_dir + with pytest.raises(Exception): # Expecting NameError or FileNotFoundError + client.get("/p/some/react/path")