api-server: first impl of rio (#960)

* first impl of rio

Signed-off-by: Teo Koon Peng <teokoonpeng@gmail.com>

* fix lint

Signed-off-by: Teo Koon Peng <teokoonpeng@gmail.com>

* cleanup

Signed-off-by: Teo Koon Peng <teokoonpeng@gmail.com>

* add option to reset app before each test

Signed-off-by: Teo Koon Peng <teokoonpeng@gmail.com>

* fix lint

Signed-off-by: Teo Koon Peng <teokoonpeng@gmail.com>

---------

Signed-off-by: Teo Koon Peng <teokoonpeng@gmail.com>
(cherry picked from commit 764c551)
Signed-off-by: Aaron Chong <aaronchongth@gmail.com>

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)]
 )

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

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]

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,

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()

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()

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 *

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)

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"])

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):