fix callbacks in MQTTHandler; enhance measurement result processing; add saving processed cells

robot-control/src/api/grbl_handler.py

@@ -70,7 +70,7 @@ class GRBLHandler:
     async def wait_until_idle(self, timeout_s):
         """Wait until GRBL reports idle status"""
         if self.debug:
-            asyncio.sleep(1)
+            await asyncio.sleep(1)
             return
             
         start = datetime.datetime.now()

robot-control/src/api/mqtt_handler.py

@@ -6,22 +6,22 @@ from utils.logging import LoggerSingleton
 
 logger = LoggerSingleton.get_logger()
 
-class Device:
+class MQTTDevice:
     def __init__(self, device_id: str, num_slots: int):
         self.device_id = device_id
         self.num_slots = num_slots
 
 class MeasurementResult(BaseModel):
-    device_id: str
     cell_id: int
-    slot: int
+    device_id: str
+    slot_id: int
     capacity: float
     status: str
 
 class MQTTHandler:
     def __init__(self, broker="localhost", port=1883, username=None, password=None):
         self.client = mqtt.Client()
-        self.devices: List[Device] = []
+        self.devices: List[MQTTDevice] = []
         self.measurement_callbacks: Dict[str, Dict[int, Callable]] = {}
         
         self.client.username_pw_set(username, password)
@@ -35,10 +35,14 @@ class MQTTHandler:
         self.client.connect(broker, port, 60)
         self.client.loop_start()
 
-    def register_device(self, device: Device):
+    def register_device(self, device_id, num_slots, callback: Callable = None):
         """Register a new device to handle"""
+        device = MQTTDevice(device_id, num_slots)
         self.devices.append(device)
-        self.measurement_callbacks[device.device_id] = {}
+        self.measurement_callbacks[device_id] = {}
+        if callback:
+            for slot in range(num_slots):
+                self.measurement_callbacks[device_id][slot] = callback
 
     def _subscribe_device_topics(self, device_id: str):
         """Subscribe to all topics for a specific device"""
@@ -69,8 +73,10 @@ class MQTTHandler:
             if topic.startswith("measurement_done/"):
                 result = MeasurementResult(**payload)
                 logger.info(f"Measurement complete {result}")
-                if device_id in self.measurement_callbacks and result.slot in self.measurement_callbacks[device_id]:
-                    self.measurement_callbacks[device_id][result.slot](result)
+                if result.device_id in self.measurement_callbacks and result.slot_id in self.measurement_callbacks[device_id]:
+                    self.measurement_callbacks[device_id][result.slot_id](result)
+                else:
+                    logger.warning(f"No callback for measurement {result}")
                     
             elif topic.startswith("soa/"):
                 logger.info(f"SOA update for device {device_id}: {payload}")
@@ -79,18 +85,13 @@ class MQTTHandler:
         except Exception as e:
             logger.error(f"Error processing message: {e}")
 
-    def start_measurement(self, device_id: str, slot: int, cell_id: str, callback: Callable = None):
+    def start_measurement(self, device_id: str, slot: int, cell_id: int):
         """Publish measurement start command for specific device"""
         if device_id not in [d.device_id for d in self.devices]:
             raise ValueError(f"Device {device_id} not registered")
             
         payload = {"slot": slot, "cell_id": cell_id}
         self.client.publish(f"cells_inserted/{device_id}", json.dumps(payload))
-        
-        if callback:
-            if device_id not in self.measurement_callbacks:
-                self.measurement_callbacks[device_id] = {}
-            self.measurement_callbacks[device_id][slot] = callback
 
     def cleanup(self):
         """Cleanup MQTT connection"""

robot-control/src/main.py

@@ -1,7 +1,7 @@
 import asyncio
 from utils.logging import LoggerSingleton
 from fastapi import FastAPI
-from robot.controller import RobotController, Cell
+from robot.controller import RobotController
 from utils.config import ConfigParser
 from vision.datamatrix import DataMatrixReader
 
@@ -32,7 +32,7 @@ class LoaderSystem:
                     self.logger.debug("No cell detected")
                     break
                 self.logger.info(f"Processing cell {cell_id}")
-                cell = Cell(cell_id)
+                cell = self.controller.add_cell(cell_id)
                 try:
                     await self.controller.pick_cell_from_feeder()
                     await self.controller.insert_cell_to_slot(cell, slot)
@@ -41,7 +41,7 @@ class LoaderSystem:
                     break
             
             # Check for completed measurements and sort cell
-            self.controller.process_finished_measurement()
+            await self.controller.process_finished_measurement()
 
 
 if __name__ == "__main__":

robot-control/src/robot/controller.py

@@ -1,10 +1,10 @@
 from dataclasses import dataclass
 from enum import Enum
 from typing import List, Tuple
-from utils.config import ConfigParser, SlotConfig
+from utils.config import ConfigParser, SlotConfig, DeviceConfig
 from robot.movement import RobotMovement
 from utils.logging import LoggerSingleton
-from api.mqtt_handler import MQTTHandler, Device
+from api.mqtt_handler import MQTTHandler, MeasurementResult
 from api.grbl_handler import GRBLHandler
 
 logger = LoggerSingleton.get_logger()
@@ -17,7 +17,7 @@ class CellStatus(Enum):
 
 @dataclass
 class Cell:
-    id: str
+    id: int
     status: CellStatus = CellStatus.WAITING
     measurement_slot: int = None
     capacity: float = None
@@ -30,6 +30,7 @@ class DropoffGrade:
 class RobotController:
     def __init__(self, config: ConfigParser):
         self.config = config
+        self.cells:dict[Cell] = {}
         self.devices = self.config.get_devices()
         self.feeder = self.config.get_feeder()
         self.dropoff_grades = self.config.get_dropoff_grades()
@@ -45,7 +46,7 @@ class RobotController:
         
         # Initialize with configured values
         self.total_slots = sum(len(device.slots) for device in self.devices)
-        self.work_queue: List[SlotConfig] = []
+        self.work_queue: List[MeasurementResult] = []
 
         self.gripper_occupied = False
         
@@ -61,7 +62,9 @@ class RobotController:
         # Register all devices with MQTT handler
         for device in self.devices:
             self.mqtt_handler.register_device(
-                Device(device_id=device.id, num_slots=len(device.slots))
+                device_id=device.id,
+                num_slots=len(device.slots),
+                callback=lambda measurement_result: self.work_queue.append(measurement_result)
             )
 
     async def connect(self):
@@ -73,13 +76,30 @@ class RobotController:
         # await self.movement.cleanup() TODO[SG]: Implement cleanup method in movement.py
         self.mqtt_handler.cleanup()
 
-    def get_device_by_id(self, device_id: str):
+    def add_cell(self, cell_id: str):
+        self.cells[cell_id] = Cell(cell_id)
+        return self.cells[cell_id]
+
+    def get_cell_by_id(self, cell_id: str) -> Cell:
+        return self.cells.get(cell_id, None)
+    
+    def get_device_by_id(self, device_id: str) -> DeviceConfig:
         for device in self.devices:
             if device.id == device_id:
                 return device
+        logger.error(f"Device {device_id} not found")
         return None
+    
+    def get_slot_by_id(self, device_id: str, slot_id: int) -> Cell:
+        try:
+            return self.get_device_by_id(device_id).slots[slot_id]
+        except AttributeError as e:
+            return None
+        except IndexError as e:
+            logger.error(f"Slot {slot_id} not found in device {device_id}")
+            return None
 
-    def get_next_free_slot(self):
+    def get_next_free_slot(self) -> SlotConfig:
         for device in self.devices:
             for slot in device.slots:
                 if not slot.occupied:
@@ -89,14 +109,34 @@ class RobotController:
         logger.warning("No free slots available")
         return None
     
-    def process_finished_measurement(self):
+    async def process_finished_measurement(self):
         if not self.work_queue:
             logger.warning("No finished measurements in queue")
             return
-        completed_slot = self.work_queue.pop(0)
-        cell_id = self.collect_cell_from_slot(completed_slot)
-        cell = Cell(cell_id)
-        self.sort_cell(cell)
+        
+        measurement_result = self.work_queue.pop(0)
+        cell = self.get_cell_by_id(measurement_result.cell_id)
+        try:
+            cell_status = CellStatus(measurement_result.status)
+        except ValueError as e:
+            logger.error(f"Measurement ({measurement_result}) could not be processed. Invalid cell status '{measurement_result.status}'")
+            return
+
+        waiting_slot = self.get_slot_by_id(measurement_result.device_id, measurement_result.slot_id)
+        if not waiting_slot:
+            logger.error(f"Measurement ({measurement_result}) could not be processed.")
+            return
+        await self.collect_cell_from_slot(waiting_slot)
+
+
+        if not cell: # Cell not found, create new
+            cell = Cell(measurement_result.cell_id, cell_status, capacity=measurement_result.capacity)
+            self.cells[measurement_result.cell_id] = cell
+        else:
+            cell.capacity = measurement_result.capacity
+            cell.status = cell_status
+
+        await self.sort_cell(cell)
 
     async def pick_cell_from_feeder(self):
         if self.gripper_occupied:
@@ -165,8 +205,7 @@ class RobotController:
         self.mqtt_handler.start_measurement(
             device_id=slot.device_id,
             slot=slot.slot_id,
-            cell_id=cell.id,
-            callback=lambda device_id, slot, cell_id, capacity: self.work_queue.append((cell_id, slot, device_id, capacity))
+            cell_id=cell.id
         )
         
         # Move back to safe height
@@ -212,14 +251,17 @@ class RobotController:
         if not self.gripper_occupied:
             logger.error("Gripper not occupied")
             return False
-        if cell.status == CellStatus.FAILED:
+        
+        if cell.id in self.cells:
+            del self.cells[cell.id] # Remove cell from our database TODO [SG]: Should we keep it for history?
+
+        if cell.status is CellStatus.FAILED:
             await self.dropoff_cell(self.rejected_grade)
             logger.info(f"Cell {cell.id} sorted to rejected grade")
             return True
         for name, grade in self.dropoff_grades.items():
             if cell.capacity >= grade.capacity_threshold:
                 await self.dropoff_cell(grade)
-                self.gripper_occupied = False
                 logger.info(f"Cell {cell.id} sorted to grade {name}")
                 return True
         logger.error(f"No suitable grade found for cell {cell.id} with capacity {cell.capacity}")

robot-control/src/utils/config.py

@@ -27,6 +27,7 @@ class FeederConfig:
 
 @dataclass
 class DropoffGradeConfig:
+    name: str
     position: Tuple[float, float, float]
     capacity_threshold: float
 
@@ -76,6 +77,7 @@ class ConfigParser:
         grades = {}
         for name, grade in self.config['dropoff_grades'].items():
             grades[name] = DropoffGradeConfig(
+                name = name,
                 position=tuple(grade['position']),
                 capacity_threshold=grade['capacity_threshold']
             )

robot-control/tests/test_api.py

@@ -1,6 +1,6 @@
 import pytest
 from fastapi.testclient import TestClient
-from api.mqtt_handler import MQTTHandler, Device, MeasurementResult, mqtt
+from api.mqtt_handler import MQTTHandler, MQTTDevice, MeasurementResult, mqtt
 from unittest.mock import Mock, patch
 
 @pytest.fixture
@@ -24,33 +24,29 @@ def mqtt_handler(mock_mqtt_client):
 
 class TestMQTTHandler:
     def test_device_registration(self, mqtt_handler:MQTTHandler):
-        device = Device("test_device", 4)
-        mqtt_handler.register_device(device)
-        assert device in mqtt_handler.devices
+        mqtt_handler.register_device("test_device", 4)
         assert "test_device" in mqtt_handler.measurement_callbacks
 
     def test_start_measurement(self, mqtt_handler:MQTTHandler, mock_mqtt_client:mqtt.Client):
-        device = Device("test_device", 4)
-        mqtt_handler.register_device(device)
-        
         callback = Mock()
-        mqtt_handler.start_measurement("test_device", 1, "cell123", callback)
+        mqtt_handler.register_device("test_device", 4, callback)
+        
+        mqtt_handler.start_measurement("test_device", 1, 123)
         
         mock_mqtt_client.publish.assert_called_once()
         assert mqtt_handler.measurement_callbacks["test_device"][1] == callback
 
     def test_measurement_callback(self, mqtt_handler:MQTTHandler):
-        device = Device("test_device", 4)
-        mqtt_handler.register_device(device)
-        
         callback = Mock()
-        mqtt_handler.start_measurement("test_device", 1, "cell123", callback)
+        mqtt_handler.register_device("test_device", 4, callback)
+        
+        mqtt_handler.start_measurement("test_device", 1, 123)
         
         # Simulate measurement complete message
         result = MeasurementResult(
+            cell_id=123,
             device_id="test_device",
-            cell_id="cell123",
-            slot=1,
+            slot_id=1,
             capacity=3000.0,
             status="complete"
         )
@@ -65,4 +61,4 @@ class TestMQTTHandler:
 
     def test_invalid_device(self, mqtt_handler:MQTTHandler):
         with pytest.raises(ValueError):
-            mqtt_handler.start_measurement("invalid_device", 1, "cell123")
+            mqtt_handler.start_measurement("invalid_device", 1, 123)

robot-control/tests/test_robot.py

@@ -3,9 +3,15 @@ from robot.movement import RobotMovement
 from robot.controller import RobotController, Cell, CellStatus
 from utils.config import ConfigParser
 
+from api.grbl_handler import GRBLHandler
+from utils.config import ConfigParser
+
 @pytest.fixture
 def robot_movement():
-    return RobotMovement()
+    config_parser = ConfigParser()
+    grbl_config = config_parser.get_grbl_config()
+    grbl_handler = GRBLHandler(grbl_config.get('port'), grbl_config.get('baudrate'))
+    return RobotMovement(grbl_handler)
 
 @pytest.fixture
 def robot_controller():
@@ -19,12 +25,13 @@ class TestRobotMovement:
         assert robot_movement.get_position() == (0.0, 0.0, 0.0)
         assert not robot_movement.is_moving
 
-    def test_set_speed(self, robot_movement: RobotMovement):
-        robot_movement.set_speed(200.0)
+    @pytest.mark.asyncio
+    async def test_set_speed(self, robot_movement: RobotMovement):
+        await robot_movement.set_speed(200.0)
         assert robot_movement.speed == 200.0
 
         with pytest.raises(ValueError):
-            robot_movement.set_speed(-100.0)
+            await robot_movement.set_speed(-100.0)
 
     @pytest.mark.asyncio
     async def test_move_to_position(self, robot_movement: RobotMovement):
@@ -81,7 +88,7 @@ class TestRobotController:
     @pytest.mark.asyncio
     async def test_sort_failed_cell(self, robot_controller: RobotController):
         failed_cell = Cell(
-            id="failed_cell",
+            id=123,
             status=CellStatus.FAILED,
             measurement_slot=1
         )