"""Test SetStimParams service functionality."""

from unittest.mock import Mock, patch

from radiens_core._generated import allegoserver_pb2, common_pb2
from radiens_core._services._allego_core import set_stim_params_proto
from radiens_core._transport._channel_pool import ChannelPool


class TestSetStimParamsProto:
    """Test set_stim_params_proto functionality."""

    def test_calls_stub_with_stim_params_request(self) -> None:
        """Test that set_stim_params_proto calls stub.SetStimParams with correct request."""
        # Arrange
        pool = Mock(spec=ChannelPool)
        channel = Mock()
        stub = Mock()
        response = common_pb2.StandardReply()

        pool.get_channel.return_value = channel
        stub.SetStimParams.return_value = response

        request = allegoserver_pb2.StimParams(
            stimShape=allegoserver_pb2.BIPHASIC,
            enabled=True,
            stimSysChanIdx=42,
        )

        # Mock the stub creation
        with patch(
            "radiens_core._generated.allegoserver_pb2_grpc.AllegoCoreStub",
            return_value=stub,
        ):
            # Act
            result = set_stim_params_proto(pool, "localhost:8080", request)

            # Assert
            pool.get_channel.assert_called_once_with("localhost:8080")
            stub.SetStimParams.assert_called_once_with(request)
            assert result is response
