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
import asyncio

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

@pytest.fixture
def robot_controller():
    config = ConfigParser().config
    config.grbl.port = "debug"
    config.mqtt.broker = "debug"  # connects to test mqtt broker
    mock_gpio = MockGPIO()
    feeder_queue = asyncio.Queue(config.feeder.max_num_cells)
    defeeder_queue = asyncio.Queue(config.defeeder.max_num_cells)
    controller = RobotController(
        config,
        gpio_handler=mock_gpio,
        i2c=None,
        vision=None,
        pump_controller=None,
        feeder_queue=feeder_queue,
        defeeder_queue=defeeder_queue
    )
    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()
        speed_mmin = 400
        await robot_movement.set_speed(speed_mmin)
        assert robot_movement.default_speed_mmmin == speed_mmin

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

    # 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=1234, status=CellStatus.WAITING)
        slot = robot_controller.get_next_free_slot()
        assert slot is not None
        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=1234,
            status=CellStatus.COMPLETED,
            health=100
        )
        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=0,
            status=CellStatus.COMPLETED,
            health=100
        )
        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=1234,
            status=CellStatus.COMPLETED,
            health=0.0
        )
        robot_controller.gripper_occupied = True
        await robot_controller.sort_cell(low_cap_cell)
        assert not robot_controller.gripper_occupied