cell_loader_ctrl/robot-control/tests/test_api.py

import pytest
from fastapi.testclient import TestClient
from api.mqtt_handler import MQTTHandler, Device, MeasurementResult, mqtt
from unittest.mock import Mock, patch

@pytest.fixture
def mock_mqtt_client():
    with patch('paho.mqtt.client.Client') as mock_client:
        client_instance = Mock()
        mock_client.return_value = client_instance
        yield client_instance

@pytest.fixture
def mqtt_handler(mock_mqtt_client):
    handler = MQTTHandler()
    yield handler
    handler.cleanup()

# @pytest.fixture
# def real_mqtt_handler():
#     handler = MQTTHandler(username="robot", password="robot")
#     yield handler
#     handler.cleanup()

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
        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)
        
        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)
        
        # Simulate measurement complete message
        result = MeasurementResult(
            device_id="test_device",
            cell_id="cell123",
            slot=1,
            capacity=3000.0,
            status="complete"
        )
        
        # Simulate MQTT message
        message = Mock()
        message.topic = "measurement_done/test_device"
        message.payload = result.json().encode()
        
        mqtt_handler.on_message(None, None, message)
        callback.assert_called_once()

    def test_invalid_device(self, mqtt_handler:MQTTHandler):
        with pytest.raises(ValueError):
            mqtt_handler.start_measurement("invalid_device", 1, "cell123")