cell_loader_ctrl/tests/test_api.py

import pytest
from robot_control.src.api.mqtt_handler import MQTTHandler, 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):
        callback = Mock()
        mqtt_handler.register_device("test_device", 4, callback)
        assert "test_device" in mqtt_handler.measurement_callbacks

    def test_start_measurement(self, mqtt_handler:MQTTHandler, mock_mqtt_client:mqtt.Client):
        callback = Mock()
        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):
        callback = Mock()
        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",
            slot_id=1,
            health=100,
            status="complete"
        )
        
        # Simulate MQTT message
        message = Mock()
        message.topic = "measurement_done/test_device"
        message.payload = result.model_dump_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, 123)