1 files changed, 13 insertions, 19 deletions
diff --git a/apirouters/tasks.py b/apirouters/tasks.py
index 02e091b..684cf60 100644
--- a/apirouters/tasks.py
+++ b/apirouters/tasks.py
@@ -2,28 +2,26 @@ from typing import Annotated
 from fastapi import APIRouter, Depends
 from pydantic import BaseModel
 
-from ..modules import task_types
-from .agents import auth_agent, Agent
+from . import agents
+from ..modules import ships
+from ..entities import task_types
+from ..entities.ship import Ship
 
 router = APIRouter()
 
 
 @router.get("/tasks")
-async def get_tasks(agent: Annotated[Agent, Depends(auth_agent)]):
-    ships = []
+async def get_tasks(agent: Annotated[agents.Agent, Depends(agents.auth_agent)]):
+    ret_list = []
     for current_ship in agent.ships.values():
-        ships.append(current_ship.get_data())
+        ret_list.append(current_ship.get_data())
 
-    return {"ships": ships}
+    return {"ships": ret_list}
 
 
 @router.get("/task/{ship_symbol}")
-async def get_tasks(ship_symbol: str, agent: Annotated[Agent, Depends(auth_agent)]):
-    for current_ship in agent.ships.values():
-        if current_ship.symbol == ship_symbol:
-            return current_ship.get_data()
-
-    return {"error": "Unknown ship symbol."}
+async def get_tasks(ship: Annotated[Ship, Depends(ships.get_ship)]):
+    return ship.get_data()
 
 
 class SetTaskBody(BaseModel):
@@ -31,10 +29,6 @@ class SetTaskBody(BaseModel):
 
 
 @router.post("/task/{ship_symbol}/set")
-async def get_tasks(ship_symbol: str, set_task: SetTaskBody, agent: Annotated[Agent, Depends(auth_agent)]):
-    for current_ship in agent.ships.values():
-        if current_ship.symbol == ship_symbol:
-            current_ship.set_task(set_task.task)
-            return current_ship.get_data()
-
-    return {"error": "Unknown ship symbol."}
+async def get_tasks(set_task: SetTaskBody, ship: Annotated[Ship, Depends(ships.get_ship)]):
+    ship.set_task(set_task.task)
+    return ship.get_data()