cell_loader_ctrl/tests/test_robot.py

import pytest
from robot_control.src.robot.movement import RobotMovement
from robot_control.src.robot.controller import RobotController, Cell, CellStatus
from robot_control.src.utils.config import ConfigParser
from robot_control.src.api.grbl_handler import GRBLHandler
from robot_control.src.api.gpio import MockGPIO

@pytest.fixture
def robot_movement():
    config = ConfigParser().config
    grbl = config.grbl
    grbl_handler = GRBLHandler(grbl.port, grbl.baudrate)
    return RobotMovement(config.movement, grbl_handler)

@pytest.fixture
def robot_controller():
    config = ConfigParser().config
    config.mqtt.broker = "debug"  # connects to test mqtt broker
    mock_gpio = MockGPIO()
    controller = RobotController(config, gpio_handler=mock_gpio)
    return controller

class TestRobotMovement:
    def test_init_position(self, robot_movement: RobotMovement):
        assert robot_movement.get_position() == (0.0, 0.0, 0.0)
        assert not robot_movement.is_moving

    @pytest.mark.asyncio
    async def test_set_speed(self, robot_movement: RobotMovement):
        await robot_movement.grbl.connect()
        await robot_movement.set_speed(200.0)
        assert robot_movement.speed == 200.0

        with pytest.raises(ValueError):
            await robot_movement.set_speed(-100.0)

    # TODO [SG]: Getting current position currently does not work in debug mode
    @pytest.mark.asyncio
    async def test_move_to_position(self, robot_movement: RobotMovement):
        await robot_movement.grbl.connect()
        result = await robot_movement.move_to_position(10.0, 20.0, 30.0)
        assert result == (10.0, 20.0, 30.0)

    @pytest.mark.asyncio
    async def test_move_relative(self, robot_movement: RobotMovement):
        pos = (100.0, 100.0, 100.0)
        rel_pos = (50.0, -50.0, 25.0)
        await robot_movement.grbl.connect()
        await robot_movement.move_to_position(*pos)
        result = await robot_movement.move_relative(*rel_pos)
        # assert result == pos + rel_pos
        # TODO [SG]: Check doesnt work for now in the debug mode

class TestRobotController:
    def test_init_controller(self, robot_controller: RobotController):
        assert len(robot_controller.devices) > 0
        assert robot_controller.feeder is not None
        assert not robot_controller.gripper_occupied

    @pytest.mark.asyncio
    async def test_get_next_free_slot(self, robot_controller: RobotController):
        slot = robot_controller.get_next_free_slot()
        assert slot is not None
        assert not slot.occupied

    @pytest.mark.asyncio
    async def test_insert_and_pick_cell(self, robot_controller: RobotController):
        # Test cell insertion
        cell = Cell(id="test_cell", status=CellStatus.WAITING)
        slot = robot_controller.get_next_free_slot()
        robot_controller.gripper_occupied = True
        
        await robot_controller.insert_cell_to_slot(cell, slot)
        assert slot.occupied
        assert slot.cell_id == cell.id

        # Test cell collection
        collected_cell = await robot_controller.pick_cell_from_slot(slot)
        assert not slot.occupied
        assert slot.cell_id is None
        assert collected_cell == cell.id

    @pytest.mark.asyncio
    async def test_sort_empty_gripper(self, robot_controller: RobotController):
        cell = Cell(
            id="test_cell",
            status=CellStatus.COMPLETED,
            capacity=3000.0
        )
        robot_controller.gripper_occupied = False
        assert not await robot_controller.sort_cell(cell)

    @pytest.mark.asyncio
    async def test_sort_failed_cell(self, robot_controller: RobotController):
        failed_cell = Cell(
            id=123,
            status=CellStatus.ERROR
        )
        robot_controller.gripper_occupied = True
        await robot_controller.sort_cell(failed_cell)
        assert not robot_controller.gripper_occupied
        
    @pytest.mark.asyncio
    async def test_sort_cell_by_capacity(self, robot_controller: RobotController):
        # Test sorting a cell with capacity that matches a grade
        cell = Cell(
            id="test_cell",
            status=CellStatus.COMPLETED,
            capacity=3000.0
        )
        robot_controller.gripper_occupied = True
        await robot_controller.sort_cell(cell)
        assert not robot_controller.gripper_occupied

        # Test sorting a cell with capacity that doesn't match any grade
        low_cap_cell = Cell(
            id="low_cap_cell",
            status=CellStatus.COMPLETED,
            capacity=0.0
        )
        robot_controller.gripper_occupied = True
        await robot_controller.sort_cell(low_cap_cell)
        assert not robot_controller.gripper_occupied