Make queries.get_severity_counts() return a dict

argos/server/queries.py

@@ -105,7 +105,7 @@ async def update_from_config(db: Session, config: schemas.Config):
     db.commit()
 
 
-async def get_severity_counts(db: Session):
+async def get_severity_counts(db: Session) -> dict:
     # Get the last result of each task
     subquery = (
         db.query(Result.task_id, func.max(Result.id).label("max_result_id"))
@@ -122,6 +122,10 @@ async def get_severity_counts(db: Session):
 
     # Execute the query and fetch the results
     task_counts_by_severity = query.all()
+
+    counts_dict = dict(task_counts_by_severity)
+    for key in ("ok", "warning", "critical"):
+        counts_dict.setdefault(key, 0)
     return task_counts_by_severity
 
 

argos/server/routes/api.py

@@ -73,8 +73,4 @@ async def get_stats(db: Session = Depends(get_db)):
 @route.get("/severities")
 async def get_severity_counts(db: Session = Depends(get_db)):
     """Returns the number of results per severity"""
-    counts = await queries.get_severity_counts(db)
-    counts_dict = dict(counts)
-    for key in ("ok", "warning", "critical"):
-        counts_dict.setdefault(key, 0)
-    return counts_dict
+    return await queries.get_severity_counts(db)

argos/server/routes/views.py

@@ -19,10 +19,7 @@ templates = Jinja2Templates(directory="argos/server/templates")
 @route.get("/")
 async def get_severity_counts(request: Request, db: Session = Depends(get_db)):
     """Returns the number of results per severity"""
-    counts = await queries.get_severity_counts(db)
-    counts_dict = dict(counts)
-    for key in ("ok", "warning", "critical"):
-        counts_dict.setdefault(key, 0)
+    counts_dict = await queries.get_severity_counts(db)
 
     agents = db.query(Result.agent_id).distinct().all()
 
@@ -35,6 +32,7 @@ async def get_severity_counts(request: Request, db: Session = Depends(get_db)):
         },
     )
 
+
 @route.get("/details")
 async def read_tasks(request: Request, db: Session = Depends(get_db)):
     tasks = db.query(Task).order_by(Task.domain).all()