From d31b7a22ab38c6d2f467f021bcfd10bd310b1217 Mon Sep 17 00:00:00 2001 From: Teo Koon Peng Date: Thu, 4 Jul 2024 11:47:08 +0800 Subject: [PATCH] api-server: first impl of rio (#960) * first impl of rio Signed-off-by: Teo Koon Peng * fix lint Signed-off-by: Teo Koon Peng * cleanup Signed-off-by: Teo Koon Peng * add option to reset app before each test Signed-off-by: Teo Koon Peng * fix lint Signed-off-by: Teo Koon Peng --------- Signed-off-by: Teo Koon Peng (cherry picked from commit 764c551353dfb6ac4fee5a3bb27241721986c021) Signed-off-by: Aaron Chong --- packages/api-server/api_server/app.py | 1 + .../api-server/api_server/models/__init__.py | 1 + packages/api-server/api_server/models/rio.py | 11 +++ .../models/tortoise_models/__init__.py | 1 + .../api_server/models/tortoise_models/rio.py | 8 +++ .../api-server/api_server/rmf_io/events.py | 8 +++ .../api-server/api_server/routes/__init__.py | 1 + packages/api-server/api_server/routes/rios.py | 44 ++++++++++++ .../api-server/api_server/routes/test_rios.py | 68 +++++++++++++++++++ .../api_server/test/test_fixtures.py | 3 +- 10 files changed, 145 insertions(+), 1 deletion(-) create mode 100644 packages/api-server/api_server/models/rio.py create mode 100644 packages/api-server/api_server/models/tortoise_models/rio.py create mode 100644 packages/api-server/api_server/routes/rios.py create mode 100644 packages/api-server/api_server/routes/test_rios.py diff --git a/packages/api-server/api_server/app.py b/packages/api-server/api_server/app.py index af72d375b..750564f56 100644 --- a/packages/api-server/api_server/app.py +++ b/packages/api-server/api_server/app.py @@ -220,6 +220,7 @@ def on_signal(sig, frame): app.include_router( routes.fleets_router, prefix="/fleets", dependencies=[Depends(user_dep)] ) +app.include_router(routes.rios_router, prefix="/rios", dependencies=[Depends(user_dep)]) app.include_router( routes.admin_router, prefix="/admin", dependencies=[Depends(user_dep)] ) diff --git a/packages/api-server/api_server/models/__init__.py b/packages/api-server/api_server/models/__init__.py index 2e3a2eba7..ac3182b97 100644 --- a/packages/api-server/api_server/models/__init__.py +++ b/packages/api-server/api_server/models/__init__.py @@ -9,6 +9,7 @@ from .labels import * from .lifts import * from .pagination import * +from .rio import * from .rmf_api.activity_discovery_request import ActivityDiscoveryRequest from .rmf_api.activity_discovery_response import ActivityDiscovery from .rmf_api.cancel_task_request import CancelTaskRequest diff --git a/packages/api-server/api_server/models/rio.py b/packages/api-server/api_server/models/rio.py new file mode 100644 index 000000000..013b9405d --- /dev/null +++ b/packages/api-server/api_server/models/rio.py @@ -0,0 +1,11 @@ +from typing import Any + +from pydantic import BaseModel, ConfigDict + + +class Rio(BaseModel): + model_config = ConfigDict(from_attributes=True) + + id: str + type: str + data: dict[str, Any] diff --git a/packages/api-server/api_server/models/tortoise_models/__init__.py b/packages/api-server/api_server/models/tortoise_models/__init__.py index 9965fdec3..ad6c7f646 100644 --- a/packages/api-server/api_server/models/tortoise_models/__init__.py +++ b/packages/api-server/api_server/models/tortoise_models/__init__.py @@ -9,6 +9,7 @@ from .ingestor_state import IngestorState from .lift_state import LiftState from .log import LogMixin +from .rio import * from .scheduled_task import * from .tasks import ( TaskEventLog, diff --git a/packages/api-server/api_server/models/tortoise_models/rio.py b/packages/api-server/api_server/models/tortoise_models/rio.py new file mode 100644 index 000000000..3a83bbbcf --- /dev/null +++ b/packages/api-server/api_server/models/tortoise_models/rio.py @@ -0,0 +1,8 @@ +import tortoise +from tortoise.fields import CharField, JSONField + + +class Rio(tortoise.Model): + id = CharField(max_length=255, pk=True) + type = CharField(max_length=255, index=True) + data = JSONField() diff --git a/packages/api-server/api_server/rmf_io/events.py b/packages/api-server/api_server/rmf_io/events.py index b6da065be..d862009bc 100644 --- a/packages/api-server/api_server/rmf_io/events.py +++ b/packages/api-server/api_server/rmf_io/events.py @@ -64,3 +64,11 @@ def __init__(self): @singleton_dep def get_beacon_events(): return BeaconEvents() + + +class RioEvents: + def __init__(self): + self.rios = Subject[mdl.Rio]() + + +rio_events = RioEvents() diff --git a/packages/api-server/api_server/routes/__init__.py b/packages/api-server/api_server/routes/__init__.py index ac3fbd34c..98eb57367 100644 --- a/packages/api-server/api_server/routes/__init__.py +++ b/packages/api-server/api_server/routes/__init__.py @@ -10,4 +10,5 @@ from .internal import router as internal_router from .lifts import router as lifts_router from .main import router as main_router +from .rios import router as rios_router from .tasks import * diff --git a/packages/api-server/api_server/routes/rios.py b/packages/api-server/api_server/routes/rios.py new file mode 100644 index 000000000..267ae4fb9 --- /dev/null +++ b/packages/api-server/api_server/routes/rios.py @@ -0,0 +1,44 @@ +from typing import Annotated + +from fastapi import Query, Response + +from api_server.fast_io import FastIORouter, SubscriptionRequest +from api_server.models import Rio +from api_server.models.tortoise_models import Rio as DbRio +from api_server.rmf_io import rio_events + +router = FastIORouter(tags=["RIOs"]) + + +@router.get("", response_model=list[Rio]) +async def query_rios( + id_: Annotated[ + str | None, Query(alias="id", description="comma separated list of ids") + ] = None, + type_: Annotated[ + str | None, Query(alias="type", description="comma separated list of types") + ] = None, +): + filters = {} + if id_: + filters["id__in"] = id_.split(",") + if type_: + filters["type__in"] = type_.split(",") + + rios = await DbRio.filter(**filters) + return [Rio.model_validate(x) for x in rios] + + +@router.sub("", response_model=Rio) +async def sub_rio(_req: SubscriptionRequest): + return rio_events.rios + + +@router.put("", response_model=None) +async def put_rio(rio: Rio, resp: Response): + rio_dict = rio.model_dump() + del rio_dict["id"] + _, created = await DbRio.update_or_create(rio_dict, id=rio.id) + if created: + resp.status_code = 201 + rio_events.rios.on_next(rio) diff --git a/packages/api-server/api_server/routes/test_rios.py b/packages/api-server/api_server/routes/test_rios.py new file mode 100644 index 000000000..e21e4b4e3 --- /dev/null +++ b/packages/api-server/api_server/routes/test_rios.py @@ -0,0 +1,68 @@ +import pydantic + +from api_server.models import Rio +from api_server.models.tortoise_models import Rio as DbRio +from api_server.rmf_io import rio_events +from api_server.test import AppFixture + + +@AppFixture.reset_app_before_test +class TestRiosRoute(AppFixture): + def test_get_rios(self): + self.portal.call( + DbRio(id="test_rio", type="test_type", data={"battery": 1}).save + ) + self.portal.call( + DbRio(id="test_rio2", type="test_type", data={"battery": 0.5}).save + ) + self.portal.call( + DbRio(id="test_rio3", type="test_type3", data={"battery": 0}).save + ) + + test_cases = [ + ("id=test_rio,test_rio2", 2), + ("id=test_rio,test_rio4", 1), + ("type=test_type,test_type3", 3), + ("type=test_type,test_rio", 2), + ("id=test_rio,test_rio3&type=test_type3", 1), + ] + + for tc in test_cases: + resp = self.client.get(f"/rios?{tc[0]}") + self.assertEqual(200, resp.status_code, tc) + rios = pydantic.TypeAdapter(list[Rio]).validate_json(resp.content) + self.assertEqual(tc[1], len(rios)) + + def test_sub_rios(self): + with self.subscribe_sio("/rios") as sub: + rio_events.rios.on_next( + Rio(id="test_rio", type="test_type", data={"battery": 1}) + ) + rio = Rio(**next(sub)) + self.assertEqual("test_rio", rio.id) + + def test_put_rios(self): + resp = self.client.put( + "/rios", + content=Rio( + id="test_rio", type="test_type", data={"battery": 1} + ).model_dump_json(), + ) + self.assertEqual(201, resp.status_code) + + rios = self.portal.call(DbRio.all) + self.assertEqual(1, len(rios)) + + resp = self.client.put( + "/rios", + content=Rio( + id="test_rio", type="test_type", data={"battery": 0.5} + ).model_dump_json(), + ) + # should return 200 if an existing resource is updated + self.assertEqual(200, resp.status_code) + rios = self.portal.call(DbRio.all) + self.assertEqual(1, len(rios)) + if not isinstance(rios[0].data, dict): + self.fail("data should be a dict") + self.assertEqual(0.5, rios[0].data["battery"]) diff --git a/packages/api-server/api_server/test/test_fixtures.py b/packages/api-server/api_server/test/test_fixtures.py index 03ce6244c..e60a73e48 100644 --- a/packages/api-server/api_server/test/test_fixtures.py +++ b/packages/api-server/api_server/test/test_fixtures.py @@ -140,6 +140,7 @@ def setUp(self): self.setUpApp() self.addCleanup(self.client.__exit__) + self.test_time = 0 self.portal = self.get_portal() @classmethod @@ -186,7 +187,7 @@ async def handle_resp(emit_room, msg, *_args, **_kwargs): if emit_room == "subscribe" and not msg["success"]: # FIXME # pylint: disable=broad-exception-raised - raise Exception("Failed to subscribe") + raise Exception("Failed to subscribe", msg) if emit_room == room: async with condition: if isinstance(msg, pydantic.BaseModel):