refactored code, use get_ship dependency for ship related calls

10 files changed, 92 insertions, 120 deletions

diff --git a/__main__.py b/__main__.py
index 5372651..662d7f3 100644
--- a/__main__.py
+++ b/__main__.py
@@ -1,15 +1,13 @@
 from contextlib import asynccontextmanager
-from typing import Annotated
 
-from fastapi import FastAPI, Depends
+from fastapi import FastAPI
 
 from .apirouters import agents, tasks, customize_ship
-from .apirouters import tasks, customize_ship
 from .modules.database import cursor, conn
 
 
 @asynccontextmanager
-async def lifespan(app: FastAPI):
+async def lifespan(application: FastAPI):
     cursor.execute(
         "CREATE TABLE IF NOT EXISTS agents(primary_key INTEGER PRIMARY KEY, callsign TEXT NOT NULL UNIQUE, token TEXT NOT NULL)")
     cursor.execute(
diff --git a/apirouters/agents.py b/apirouters/agents.py
index 1b927e7..49d3740 100644
--- a/apirouters/agents.py
+++ b/apirouters/agents.py
@@ -1,30 +1,17 @@
 from enum import Enum
 from typing import Annotated, Dict
 
-import openapi_client
 from fastapi import APIRouter, Depends, Response, HTTPException
 from fastapi.security import OAuth2PasswordBearer, OAuth2PasswordRequestForm
 
-from ..modules import task_types
 from ..modules.database import cursor, conn
-from ..modules.ships import Ship
+from ..entities import task_types
+from ..entities.agent import Agent
+from ..entities.ship import Ship
 
 router = APIRouter()
 oauth2_scheme = OAuth2PasswordBearer(tokenUrl="token")
-agents: Dict[str, 'Agent'] = {}
-
-
-class Agent:
-    def __init__(self, agent_symbol, token):
-        config = openapi_client.Configuration(access_token=token)
-        self.api_client = openapi_client.ApiClient(config)
-
-        self.agents_api = openapi_client.AgentsApi(self.api_client)
-        self.fleet_api = openapi_client.FleetApi(self.api_client)
-        self.systems_api = openapi_client.SystemsApi(self.api_client)
-
-        self.agent_symbol = agent_symbol
-        self.ships: Dict[str, Ship] = {}
+agents: Dict[str, Agent] = {}
 
 
 @router.post("/token")
diff --git a/apirouters/customize_ship.py b/apirouters/customize_ship.py
index c69d421..ca3ef83 100644
--- a/apirouters/customize_ship.py
+++ b/apirouters/customize_ship.py
@@ -2,7 +2,8 @@ from typing import Annotated
 from fastapi import APIRouter, Depends
 from pydantic import BaseModel
 
-from .agents import auth_agent, Agent
+from ..modules import ships
+from ..entities.ship import Ship
 
 router = APIRouter()
 
@@ -11,12 +12,8 @@ class RenameBody(BaseModel):
     name: str
 
 
-# TODO depend on ship instead of auth_agent
 @router.post("/customize_ship/{ship_symbol}/rename")
-async def rename(ship_symbol: str, rename_body: RenameBody, agent: Annotated[Agent, Depends(auth_agent)]):
-    for current_ship in agent.ships.values():
-        if current_ship.symbol == ship_symbol:
-            current_ship.rename(rename_body.name)
-            return current_ship.get_data()
+async def rename(ship_symbol: str, rename_body: RenameBody, ship: Annotated[Ship, Depends(ships.get_ship)]):
+    ship.rename(rename_body.name)
+    return ship.get_data()
 
-    return {"error": "Unknown ship symbol."}
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()
diff --git a/datatemplates/__init__.py b/entities/__init__.py
index e69de29..e69de29 100644
--- a/datatemplates/__init__.py
+++ b/entities/__init__.py
diff --git a/entities/agent.py b/entities/agent.py
new file mode 100644
index 0000000..ba99dac
--- /dev/null
+++ b/entities/agent.py
@@ -0,0 +1,17 @@
+import openapi_client
+from typing import Dict
+
+from .ship import Ship
+
+
+class Agent:
+    def __init__(self, agent_symbol, token):
+        config = openapi_client.Configuration(access_token=token)
+        self.api_client = openapi_client.ApiClient(config)
+
+        self.agents_api = openapi_client.AgentsApi(self.api_client)
+        self.fleet_api = openapi_client.FleetApi(self.api_client)
+        self.systems_api = openapi_client.SystemsApi(self.api_client)
+
+        self.agent_symbol = agent_symbol
+        self.ships: Dict[str, Ship] = {}
diff --git a/entities/ship.py b/entities/ship.py
new file mode 100644
index 0000000..922f961
--- /dev/null
+++ b/entities/ship.py
@@ -0,0 +1,41 @@
+from typing import Dict
+
+from . import task_types
+from ..modules.database import cursor, conn
+
+
+class Ship:
+    def __init__(self, symbol: str):
+        self.symbol = symbol
+        self.callsign = "-".split(symbol)[0]
+
+        self.task = None
+        self.name = None
+
+    def get_data(self) -> Dict:
+        return {"name": self.name, "symbol": self.symbol, "task": self.task}
+
+    def load_task(self):
+        cursor.execute("SELECT name, task, params FROM ships WHERE symbol = ?", (self.symbol,))
+        row = cursor.fetchone()
+        if row is None:
+            cursor.execute("INSERT INTO ships (symbol, task, params, name) VALUES (?, ?, ?, ?)",
+                           (self.symbol, task_types.IDLE, None, self.symbol))
+            conn.commit()
+
+            self.task = task_types.IDLE
+            self.task = self.symbol
+        else:
+            self.name = row[0]
+            self.task = row[1]
+
+    def set_task(self, task):
+        self.task = task
+        if task != task_types.ERROR:
+            cursor.execute("UPDATE ships SET task = ? WHERE symbol = ?", (task, self.symbol))
+            conn.commit()
+
+    def rename(self, name):
+        self.name = name
+        cursor.execute("UPDATE ships SET name = ? WHERE symbol = ?", (name, self.symbol))
+        conn.commit()
diff --git a/modules/task_types.py b/entities/task_types.py
index d23ce55..d23ce55 100644
--- a/modules/task_types.py
+++ b/entities/task_types.py
diff --git a/modules/daemon.py b/modules/daemon.py
deleted file mode 100644
index 085a7f6..0000000
--- a/modules/daemon.py
+++ /dev/null
@@ -1,36 +0,0 @@
-import openapi_client
-from fastapi import FastAPI
-from pprint import pprint
-
-fast_api = FastAPI()
-
-class Daemon:
-    def __init__(self, agent_symbol, token):
-        config = openapi_client.Configuration(access_token=token)
-        self.api_client = openapi_client.ApiClient(config)
-
-        self.agents_api = openapi_client.AgentsApi(self.api_client)
-        self.fleet_api = openapi_client.FleetApi(self.api_client)
-        self.systems_api = openapi_client.SystemsApi(self.api_client)
-
-        self.agent_symbol = agent_symbol
-
-    @fast_api.get("/")
-    def foo(self):
-        return "hello"
-
-    def run(self):
-        try:
-            api_response = self.agents_api.get_agent(self.agent_symbol)
-            print("The response of AgentsApi->get_agent:\n")
-            print(api_response)
-
-            api_response = self.fleet_api.get_my_ships()
-            for elem in api_response.data:
-                print(elem.symbol, elem.nav.waypoint_symbol, elem.nav.status)
-                waypoint = self.systems_api.get_waypoint(elem.nav.system_symbol, elem.nav.waypoint_symbol).data
-
-                print(waypoint.symbol, waypoint.type)
-
-        except Exception as e:
-            print("Exception when calling AgentsApi->get_agent: %s\n" % e)
diff --git a/modules/ships.py b/modules/ships.py
index 05ce6c1..9b0d4bd 100644
--- a/modules/ships.py
+++ b/modules/ships.py
@@ -1,41 +1,15 @@
-from typing import Dict
+from typing import Annotated
 
-from . import task_types
-from .database import cursor, conn
+from fastapi import Depends, HTTPException
 
+from ..apirouters import agents
+from ..entities.agent import Agent
+from ..entities.ship import Ship
 
-class Ship:
-    def __init__(self, symbol: str):
-        self.symbol = symbol
-        self.callsign = "-".split(symbol)[0]
 
-        self.task = None
-        self.name = None
+async def get_ship(ship_symbol: str, agent: Annotated[Agent, Depends(agents.auth_agent)]) -> Ship:
+    for ship in agent.ships.values():
+        if ship.symbol == ship_symbol:
+            return ship
 
-    def get_data(self) -> Dict:
-        return {"name": self.name, "symbol": self.symbol, "task": self.task}
-
-    def load_task(self):
-        cursor.execute("SELECT name, task, params FROM ships WHERE symbol = ?", (self.symbol,))
-        row = cursor.fetchone()
-        if row is None:
-            cursor.execute("INSERT INTO ships (symbol, task, params, name) VALUES (?, ?, ?, ?)",
-                           (self.symbol, task_types.IDLE, None, self.symbol))
-            conn.commit()
-
-            self.task = task_types.IDLE
-            self.task = self.symbol
-        else:
-            self.name = row[0]
-            self.task = row[1]
-
-    def set_task(self, task):
-        self.task = task
-        if task != task_types.ERROR:
-            cursor.execute("UPDATE ships SET task = ? WHERE symbol = ?", (task, self.symbol))
-            conn.commit()
-
-    def rename(self, name):
-        self.name = name
-        cursor.execute("UPDATE ships SET name = ? WHERE symbol = ?", (name, self.symbol))
-        conn.commit()
+    raise HTTPException(status_code=400, detail="Unknown ship symbol.")