Fix html serving. /setup now returns /setup.html. Svelte is happy.

Fix tests: create build files they expect to be there, use a mock directory
Kiln-AI · Oct 1, 2024 · d117a1a · d117a1a
1 parent 1c6295b
commit d117a1a
Show file tree

Hide file tree

Showing 2 changed files with 117 additions and 3 deletions.
diff --git a/libs/studio/kiln_studio/server.py b/libs/studio/kiln_studio/server.py
@@ -106,7 +106,24 @@ def read_item(item_id: int, q: Union[str, None] = None):
 
 
 # Web UI
-app.mount("/", StaticFiles(directory=studio_path(), html=True), name="studio")
+# File server that maps /foo/bar to /foo/bar.html (Starlette StaticFiles only does index.html)
+class HTMLStaticFiles(StaticFiles):
+    async def get_response(self, path: str, scope):
+        try:
+            response = await super().get_response(path, scope)
+            return response
+        except Exception as e:
+            # catching HTTPException explicitly not working for some reason
+            if getattr(e, "status_code", None) == 404:
+                # Return the .html version of the file if the .html version exists
+                return await super().get_response(f"{path}.html", scope)
+            raise e
+
+
+# Ensure studio_path exists (test servers don't necessarily create it)
+os.makedirs(studio_path(), exist_ok=True)
+# Serves the web UI at root
+app.mount("/", HTMLStaticFiles(directory=studio_path(), html=True), name="studio")
 
 if __name__ == "__main__":
     uvicorn.run(app, host="127.0.0.1", port=8757)
diff --git a/libs/studio/kiln_studio/test_server.py b/libs/studio/kiln_studio/test_server.py
@@ -1,8 +1,12 @@
-from unittest.mock import patch
+import os
+import tempfile
+from unittest.mock import MagicMock, patch
 
 import pytest
 import requests
+from fastapi import HTTPException
 from fastapi.testclient import TestClient
+from kiln_studio.server import HTMLStaticFiles, studio_path
 
 from libs.studio.kiln_studio.server import app
 
@@ -92,7 +96,100 @@ def test_cors_blocked_origins(origin):
     assert "access-control-allow-origin" not in response.headers
 
 
-def test_cors_no_origin():
+@pytest.fixture
+def mock_studio_path():
+    with tempfile.TemporaryDirectory() as temp_dir:
+        with patch("kiln_studio.server.studio_path", return_value=temp_dir):
+            yield temp_dir
+
+
+def create_studio_test_file(relative_path):
+    full_path = os.path.join(studio_path(), relative_path)
+    os.makedirs(os.path.dirname(full_path), exist_ok=True)
+    with open(full_path, "w") as f:
+        f.write("<html><body>Test</body></html>")
+    return full_path
+
+
+def test_cors_no_origin(mock_studio_path):
+    # Create index.html in the mock studio path
+    create_studio_test_file("index.html")
+
+    # Use the client to get the root path
     response = client.get("/")
+
+    # Assert the response
     assert response.status_code == 200
     assert "access-control-allow-origin" not in response.headers
+
+
+class TestHTMLStaticFiles:
+    @pytest.fixture
+    def html_static_files(self):
+        import os
+        import tempfile
+
+        self.test_dir = tempfile.mkdtemp()
+        with open(os.path.join(self.test_dir, "existing_file"), "w") as f:
+            f.write("Test content")
+        return HTMLStaticFiles(directory=self.test_dir, html=True)
+
+    @pytest.mark.asyncio
+    async def test_get_response_existing_file(self, html_static_files):
+        with patch("fastapi.staticfiles.StaticFiles.get_response") as mock_get_response:
+            mock_response = MagicMock()
+            mock_get_response.return_value = mock_response
+
+            response = await html_static_files.get_response("existing_file", {})
+
+            assert response == mock_response
+            mock_get_response.assert_called_once_with("existing_file", {})
+
+    @pytest.mark.asyncio
+    async def test_get_response_html_fallback(self, html_static_files):
+        with patch("fastapi.staticfiles.StaticFiles.get_response") as mock_get_response:
+
+            def side_effect(path, scope):
+                if path.endswith(".html"):
+                    return MagicMock()
+                raise HTTPException(status_code=404)
+
+            mock_get_response.side_effect = side_effect
+
+            response = await html_static_files.get_response("non_existing_file", {})
+
+            assert response is not None
+            assert mock_get_response.call_count == 2
+            mock_get_response.assert_any_call("non_existing_file", {})
+            mock_get_response.assert_any_call("non_existing_file.html", {})
+
+    @pytest.mark.asyncio
+    async def test_get_response_not_found(self, html_static_files):
+        with patch("fastapi.staticfiles.StaticFiles.get_response") as mock_get_response:
+            mock_get_response.side_effect = HTTPException(status_code=404)
+
+            with pytest.raises(HTTPException):
+                await html_static_files.get_response("non_existing_file", {})
+
+    @pytest.mark.asyncio
+    async def test_setup_route(self, mock_studio_path):
+        import os
+
+        # Ensure studio_path exists
+        os.makedirs(studio_path(), exist_ok=True)
+        create_studio_test_file("index.html")
+        create_studio_test_file("setup.html")
+        create_studio_test_file("setup/connect_providers/index.html")
+
+        # root index.html
+        response = client.get("/")
+        assert response.status_code == 200
+        # setup.html
+        response = client.get("/setup")
+        assert response.status_code == 200
+        # nested index.html
+        response = client.get("/setup/connect_providers")
+        assert response.status_code == 200
+        # non existing file
+        response = client.get("/setup/non_existing_file")
+        assert response.status_code == 404