From acd3c25079593666a5c261550289f76a2a695e27 Mon Sep 17 00:00:00 2001
From: WangSL <48207171+WSL0809@users.noreply.github.com>
Date: Fri, 12 Apr 2024 15:32:20 +0800
Subject: [PATCH] manage_plan.py

---
 api/manage_plan.py | 21 ++++++++++++---------
 1 file changed, 12 insertions(+), 9 deletions(-)

diff --git a/api/manage_plan.py b/api/manage_plan.py
index dbfaf30..1b1b5f6 100644
--- a/api/manage_plan.py
+++ b/api/manage_plan.py
@@ -1,3 +1,5 @@
+from typing import List
+
 from fastapi import APIRouter, Depends, HTTPException
 from pydantic import BaseModel
 from sqlalchemy import text
@@ -24,6 +26,11 @@ class Config:
         orm_mode = True
 
 
+class PlanResp(BaseModel):
+    meal_plan: List[PlanRecvBase]
+    recovery_plan: List[PlanRecvBase]
+
+
 def create_plan(db, plan_recv):
     create_plan_sql = ''
     if plan_recv.plan_category == "meal_plan":
@@ -66,7 +73,7 @@ def delete_plan(db, plan_recv):
         raise HTTPException(status_code=500, detail=str(e))
 
 
-def get_plans(db, plan_category) -> list[dict]:
+def get_plans(db, plan_category):
     get_plans_sql = ''
     if plan_category == "meal_plan":
         get_plans_sql = text(
@@ -82,7 +89,8 @@ def get_plans(db, plan_category) -> list[dict]:
     except SQLAlchemyError as e:
         db.rollback()
         raise HTTPException(status_code=500, detail=str(e))
-    return [dict(row) for row in res]
+    return [PlanRecvBase(**dict(row)) for row in res]
+
 
 # add new plan
 @router.post("/new_plan")
@@ -118,15 +126,10 @@ async def delete_plan(plan_recv: PlanRecv, current_user: User = Depends(get_curr
     }
 
 
-@router.get("/get_plans")
+@router.get("/get_plans", response_model=PlanResp)
 async def get_plans(current_user: User = Depends(get_current_active_user), db: Session = Depends(get_db)):
     if current_user.role != "admin":
         raise HTTPException(status_code=401, detail="only admin can get plans")
     meal_plan_res = get_plans(db, "meal_plan")
     recovery_plan_res = get_plans(db, "recovery_plan")
-    return {
-        "meal_plan": meal_plan_res,
-        "recovery_plan": recovery_plan_res,
-        "status": "success",
-        "details": "get plans success"
-    }
+    return PlanResp(meal_plan=meal_plan_res, recovery_plan=recovery_plan_res)