Remove flask from mx bluesky

pyproject.toml

@@ -28,7 +28,7 @@ dependencies = [
     "annotated_types",
     "caproto",
     "fastapi[all]",
-    "flask-restful",
+    "uvicorn",
     "jupyterlab",
     "matplotlib",
     "nexgen >= 0.11.0",

src/mx_bluesky/hyperion/plan_runner_api.py

@@ -1,55 +1,45 @@
-from threading import Thread
+from typing import Annotated
 
-from flask import Flask
-from flask_restful import Api, Resource
+import uvicorn
+from fastapi import Depends, FastAPI
 
 from mx_bluesky.common.utils.log import LOGGER
 from mx_bluesky.hyperion.plan_runner import PlanRunner
 
+app = FastAPI()
 
-# Ignore this function for code coverage as there is no way to shut down
-# a server once it is started.
-def create_server_for_udc(runner: PlanRunner, port: int) -> Thread:  # pragma: no cover
-    """Create a minimal API for Hyperion UDC mode"""
-    app = create_app_for_udc(runner)
-
-    flask_thread = Thread(
-        target=app.run,
-        kwargs={"host": "0.0.0.0", "port": port},
-        daemon=True,
-    )
-    flask_thread.start()
-    LOGGER.info(f"Hyperion now listening on {port}")
-    return flask_thread
+_plan_runner: PlanRunner
 
 
-def create_app_for_udc(runner: PlanRunner):
-    app = Flask(__name__)
-    api = Api(app)
-    api.add_resource(StatusResource, "/status", resource_class_args=[runner])
-    api.add_resource(CallbackLiveness, "/callbackPing", resource_class_args=[runner])
-    return app
+# Ignore this function for code coverage as there is no way to shut down
+# a server once it is started.
+def create_server_for_udc(
+    runner: PlanRunner, port: int
+) -> uvicorn.Server:  # pragma: no cover
+    # register resources with the app via instantiation
+    global _plan_runner
+    _plan_runner = runner
 
+    """Create a minimal API for Hyperion UDC mode"""
+    config = uvicorn.Config("mx_bluesky.hyperion.plan_runner_api:app", port=port)
+    server = uvicorn.Server(config)
+    server.run()
 
-class StatusResource(Resource):
-    """Status endpoint, used by k8s healthcheck probe"""
+    LOGGER.info(f"Hyperion now listening on {port}")
+    return server
 
-    def __init__(self, runner: PlanRunner):
-        super().__init__()
-        self._runner = runner
 
-    def get(self):
-        status = self._runner.current_status
-        return {"status": status.value}
+def plan_runner() -> PlanRunner:
+    return _plan_runner
 
 
-class CallbackLiveness(Resource):
-    """Called periodically by the external callbacks to indicate that they are still running"""
+@app.get("/status")
+async def get_status(runner: Annotated[PlanRunner, Depends(plan_runner)]):
+    status = runner.current_status
+    return {"status": status.value}
 
-    def __init__(self, runner: PlanRunner):
-        super().__init__()
-        self._runner = runner
 
-    def get(self):
-        LOGGER.debug("External callback ping received.")
-        self._runner.reset_callback_watchdog_timer()
+@app.get("/callbackPing")
+async def get_callback_ping(runner: Annotated[PlanRunner, Depends(plan_runner)]):
+    LOGGER.debug("External callback ping received.")
+    runner.reset_callback_watchdog_timer()

tests/unit_tests/hyperion/test_plan_runner_api.py

@@ -1,42 +1,41 @@
-from unittest.mock import MagicMock
+from unittest.mock import MagicMock, patch
 
 import pytest
-from flask import Flask
-from flask.testing import FlaskClient
+from fastapi import FastAPI
+from fastapi.testclient import TestClient
 
 from mx_bluesky.common.parameters.constants import Status
 from mx_bluesky.hyperion.in_process_runner import InProcessRunner
-from mx_bluesky.hyperion.plan_runner_api import (
-    create_app_for_udc,
-)
+from mx_bluesky.hyperion.plan_runner_api import app
 
 
 @pytest.fixture()
 def mock_runner():
-    return MagicMock(spec=InProcessRunner)
+    runner = MagicMock(spec=InProcessRunner)
+    with patch("mx_bluesky.hyperion.plan_runner_api._plan_runner", runner, create=True):
+        yield runner
 
 
 @pytest.fixture()
 def app_under_test(mock_runner):
-    app = create_app_for_udc(mock_runner)
     yield app
 
 
 @pytest.fixture()
-def client(app_under_test: Flask) -> FlaskClient:
-    return app_under_test.test_client()
+def client(app_under_test: FastAPI) -> TestClient:
+    return TestClient(app_under_test)
 
 
 def test_plan_runner_api_fetch_status(app_under_test, client, mock_runner):
     mock_runner.current_status = Status.BUSY
     response = client.get("/status")
     assert response.status_code == 200
-    assert response.content_type == "application/json"
-    assert response.json["status"] == Status.BUSY.value
+    assert response.headers["Content-Type"] == "application/json"
+    assert response.json()["status"] == Status.BUSY.value
 
 
 def test_plan_runner_api_callback_liveness(app_under_test, client, mock_runner):
     response = client.get("/callbackPing")
     mock_runner.reset_callback_watchdog_timer.assert_called_once()
     assert response.status_code == 200
-    assert response.content_type == "application/json"
+    assert response.headers["Content-Type"] == "application/json"