fix: unit tests fixes and ruff formating

N25B-205
This commit is contained in:
Björn Otgaar
2025-10-28 11:31:05 +01:00
parent 52faa59184
commit 47a87d0b4a
11 changed files with 307 additions and 209 deletions

View File

@@ -10,13 +10,22 @@ from control_backend.schemas.ri_message import SpeechCommand
logger = logging.getLogger(__name__)
class RICommandAgent(Agent):
subsocket: zmq.Socket
pubsocket: zmq.Socket
address = ""
bind = False
def __init__(self, jid: str, password: str, port: int = 5222, verify_security: bool = False, address = "tcp://localhost:0000", bind = False):
def __init__(
self,
jid: str,
password: str,
port: int = 5222,
verify_security: bool = False,
address="tcp://localhost:0000",
bind=False,
):
super().__init__(jid, password, port, verify_security)
self.address = address
self.bind = bind
@@ -29,12 +38,12 @@ class RICommandAgent(Agent):
assert self.agent is not None
# Get a message internally (with topic command)
topic, body = await self.agent.subsocket.recv_multipart()
# Try to get body
try:
body = json.loads(body)
message = SpeechCommand.model_validate(body)
# Send to the robot.
await self.agent.pubsocket.send_json(message.model_dump())
except Exception as e:
@@ -48,11 +57,11 @@ class RICommandAgent(Agent):
# To the robot
self.pubsocket = context.socket(zmq.PUB)
if self.bind:
if self.bind:
self.pubsocket.bind(self.address)
else :
else:
self.pubsocket.connect(self.address)
# Receive internal topics regarding commands
self.subsocket = context.socket(zmq.SUB)
self.subsocket.connect(settings.zmq_settings.internal_comm_address)

View File

@@ -12,14 +12,21 @@ from control_backend.agents.ri_command_agent import RICommandAgent
logger = logging.getLogger(__name__)
class RICommunicationAgent(Agent):
req_socket: zmq.Socket
_address = ""
_bind = True
def __init__(self, jid: str, password: str, port: int = 5222,
verify_security: bool = False, address = "tcp://localhost:0000",
bind = False):
def __init__(
self,
jid: str,
password: str,
port: int = 5222,
verify_security: bool = False,
address="tcp://localhost:0000",
bind=False,
):
super().__init__(jid, password, port, verify_security)
self._address = address
self._bind = bind
@@ -37,28 +44,26 @@ class RICommunicationAgent(Agent):
# Wait up to three seconds for a reply:)
try:
message = await asyncio.wait_for(
self.agent.req_socket.recv_json(),
timeout=3.0)
message = await asyncio.wait_for(self.agent.req_socket.recv_json(), timeout=3.0)
# We didnt get a reply :(
except asyncio.TimeoutError as e:
except asyncio.TimeoutError as e:
logger.info("No ping retrieved in 3 seconds, killing myself.")
self.kill()
logger.debug("Received message \"%s\"", message)
logger.debug('Received message "%s"', message)
if "endpoint" not in message:
logger.error("No received endpoint in message, excepted ping endpoint.")
return
# See what endpoint we received
match message["endpoint"]:
case "ping":
case "ping":
await asyncio.sleep(1)
case _:
logger.info("Received message with topic different than ping," \
" while ping expected.")
logger.info(
"Received message with topic different than ping, while ping expected."
)
async def setup(self, max_retries: int = 5):
"""
@@ -67,14 +72,13 @@ class RICommunicationAgent(Agent):
logger.info("Setting up %s", self.jid)
retries = 0
# Let's try a certain amount of times before failing connection
while retries < max_retries:
# Bind request socket
self.req_socket = context.socket(zmq.REQ)
if self._bind:
if self._bind:
self.req_socket.bind(self._address)
else:
else:
self.req_socket.connect(self._address)
# Send our message and receive one back:)
@@ -85,10 +89,13 @@ class RICommunicationAgent(Agent):
received_message = await asyncio.wait_for(self.req_socket.recv_json(), timeout=20.0)
except asyncio.TimeoutError:
logger.warning("No connection established in 20 seconds (attempt %d/%d)",
retries + 1, max_retries)
logger.warning(
"No connection established in 20 seconds (attempt %d/%d)",
retries + 1,
max_retries,
)
retries += 1
continue
continue
except Exception as e:
logger.error("Unexpected error during negotiation: %s", e)
@@ -99,10 +106,14 @@ class RICommunicationAgent(Agent):
endpoint = received_message.get("endpoint")
if endpoint != "negotiate/ports":
# TODO: Should this send a message back?
logger.error("Invalid endpoint '%s' received (attempt %d/%d)",
endpoint, retries + 1, max_retries)
logger.error(
"Invalid endpoint '%s' received (attempt %d/%d)",
endpoint,
retries + 1,
max_retries,
)
retries += 1
continue
continue
# At this point, we have a valid response
try:
@@ -113,7 +124,7 @@ class RICommunicationAgent(Agent):
if not bind:
addr = f"tcp://localhost:{port}"
else:
else:
addr = f"tcp://*:{port}"
match id:
@@ -125,11 +136,13 @@ class RICommunicationAgent(Agent):
self.req_socket.bind(addr)
case "actuation":
ri_commands_agent = RICommandAgent(
settings.agent_settings.ri_command_agent_name +
'@' + settings.agent_settings.host,
settings.agent_settings.ri_command_agent_name,
address=addr,
bind=bind )
settings.agent_settings.ri_command_agent_name
+ "@"
+ settings.agent_settings.host,
settings.agent_settings.ri_command_agent_name,
address=addr,
bind=bind,
)
await ri_commands_agent.start()
case _:
logger.warning("Unhandled negotiation id: %s", id)
@@ -150,5 +163,3 @@ class RICommunicationAgent(Agent):
listen_behaviour = self.ListenBehaviour()
self.add_behaviour(listen_behaviour)
logger.info("Finished setting up %s", self.jid)

View File

@@ -9,6 +9,7 @@ logger = logging.getLogger(__name__)
router = APIRouter()
@router.post("/command", status_code=202)
async def receive_command(command: SpeechCommand, request: Request):
# Validate and retrieve data.
@@ -16,5 +17,6 @@ async def receive_command(command: SpeechCommand, request: Request):
topic = b"command"
pub_socket: Socket = request.app.state.internal_comm_socket
pub_socket.send_multipart([topic, command.model_dump_json().encode()])
return {"status": "Command received"}

View File

@@ -6,12 +6,6 @@ api_router = APIRouter()
api_router.include_router(message.router, tags=["Messages"])
api_router.include_router(
sse.router,
tags=["SSE"]
)
api_router.include_router(sse.router, tags=["SSE"])
api_router.include_router(
command.router,
tags=["Commands"]
)
api_router.include_router(command.router, tags=["Commands"])

View File

@@ -15,6 +15,7 @@ class AgentSettings(BaseModel):
ri_communication_agent_name: str = "ri_communication_agent"
ri_command_agent_name: str = "ri_command_agent"
class Settings(BaseSettings):
app_title: str = "PepperPlus"

View File

@@ -31,16 +31,19 @@ async def lifespan(app: FastAPI):
logger.info("Internal publishing socket bound to %s", internal_comm_socket)
# Initiate agents
ri_communication_agent = RICommunicationAgent(settings.agent_settings.ri_communication_agent_name +
'@' + settings.agent_settings.host,
settings.agent_settings.ri_communication_agent_name,
address="tcp://*:5555", bind=True)
ri_communication_agent = RICommunicationAgent(
settings.agent_settings.ri_communication_agent_name + "@" + settings.agent_settings.host,
settings.agent_settings.ri_communication_agent_name,
address="tcp://*:5555",
bind=True,
)
await ri_communication_agent.start()
bdi_core = BDICoreAgent(settings.agent_settings.bdi_core_agent_name +
'@' + settings.agent_settings.host,
settings.agent_settings.bdi_core_agent_name,
"src/control_backend/agents/bdi/rules.asl")
bdi_core = BDICoreAgent(
settings.agent_settings.bdi_core_agent_name + "@" + settings.agent_settings.host,
settings.agent_settings.bdi_core_agent_name,
"src/control_backend/agents/bdi/rules.asl",
)
await bdi_core.start()
yield

View File

@@ -12,7 +12,7 @@ class RIEndpoint(str, Enum):
class RIMessage(BaseModel):
endpoint: RIEndpoint
data: Any
data: Any
class SpeechCommand(RIMessage):