nix-builder-autoscaler/agent/nix_builder_autoscaler/tests/test_haproxy_provider.py

"""Unit tests for the HAProxy provider, mocking at socket level."""

from unittest.mock import MagicMock, patch

import pytest

from nix_builder_autoscaler.providers.haproxy import HAProxyError, HAProxyRuntime

# HAProxy `show stat` CSV — trimmed to columns the parser uses.
# Full output has many more columns; we keep through `status` (col 17).
SHOW_STAT_CSV = (
    "# pxname,svname,qcur,qmax,scur,smax,slim,stot,"
    "bin,bout,dreq,dresp,ereq,econ,eresp,wretr,wredis,status\n"
    "all,BACKEND,0,0,2,5,200,100,5000,6000,0,0,,0,0,0,0,UP\n"
    "all,slot001,0,0,1,3,50,50,2500,3000,0,0,,0,0,0,0,UP\n"
    "all,slot002,0,0,1,2,50,50,2500,3000,0,0,,0,0,0,0,DOWN\n"
    "all,slot003,0,0,0,0,50,0,0,0,0,0,,0,0,0,0,MAINT\n"
)


class TestSetSlotAddr:
    @patch("nix_builder_autoscaler.providers.haproxy.socket.socket")
    def test_sends_correct_command(self, mock_socket_cls):
        mock_sock = MagicMock()
        mock_socket_cls.return_value = mock_sock
        mock_sock.recv.side_effect = [b"\n", b""]

        h = HAProxyRuntime("/tmp/test.sock", "all", "slot")
        h.set_slot_addr("slot001", "100.64.0.1", 22)

        mock_sock.connect.assert_called_once_with("/tmp/test.sock")
        mock_sock.sendall.assert_called_once_with(
            b"set server all/slot001 addr 100.64.0.1 port 22\n"
        )


class TestEnableSlot:
    @patch("nix_builder_autoscaler.providers.haproxy.socket.socket")
    def test_sends_correct_command(self, mock_socket_cls):
        mock_sock = MagicMock()
        mock_socket_cls.return_value = mock_sock
        mock_sock.recv.side_effect = [b"\n", b""]

        h = HAProxyRuntime("/tmp/test.sock", "all", "slot")
        h.enable_slot("slot001")

        mock_sock.sendall.assert_called_once_with(b"enable server all/slot001\n")


class TestDisableSlot:
    @patch("nix_builder_autoscaler.providers.haproxy.socket.socket")
    def test_sends_correct_command(self, mock_socket_cls):
        mock_sock = MagicMock()
        mock_socket_cls.return_value = mock_sock
        mock_sock.recv.side_effect = [b"\n", b""]

        h = HAProxyRuntime("/tmp/test.sock", "all", "slot")
        h.disable_slot("slot001")

        mock_sock.sendall.assert_called_once_with(b"disable server all/slot001\n")


class TestReadSlotHealth:
    @patch("nix_builder_autoscaler.providers.haproxy.socket.socket")
    def test_parses_csv_correctly(self, mock_socket_cls):
        mock_sock = MagicMock()
        mock_socket_cls.return_value = mock_sock
        mock_sock.recv.side_effect = [SHOW_STAT_CSV.encode(), b""]

        h = HAProxyRuntime("/tmp/test.sock", "all", "slot")
        health = h.read_slot_health()

        assert len(health) == 3
        # BACKEND row should be excluded (svname "BACKEND" doesn't start with "slot")

        assert health["slot001"].status == "UP"
        assert health["slot001"].scur == 1
        assert health["slot001"].qcur == 0

        assert health["slot002"].status == "DOWN"
        assert health["slot002"].scur == 1
        assert health["slot002"].qcur == 0

        assert health["slot003"].status == "MAINT"
        assert health["slot003"].scur == 0
        assert health["slot003"].qcur == 0


class TestSlotIsUp:
    @patch("nix_builder_autoscaler.providers.haproxy.socket.socket")
    def test_returns_true_for_up_slot(self, mock_socket_cls):
        mock_sock = MagicMock()
        mock_socket_cls.return_value = mock_sock
        mock_sock.recv.side_effect = [SHOW_STAT_CSV.encode(), b""]

        h = HAProxyRuntime("/tmp/test.sock", "all", "slot")
        assert h.slot_is_up("slot001") is True

    @patch("nix_builder_autoscaler.providers.haproxy.socket.socket")
    def test_returns_false_for_down_slot(self, mock_socket_cls):
        mock_sock = MagicMock()
        mock_socket_cls.return_value = mock_sock
        mock_sock.recv.side_effect = [SHOW_STAT_CSV.encode(), b""]

        h = HAProxyRuntime("/tmp/test.sock", "all", "slot")
        assert h.slot_is_up("slot002") is False


class TestErrorHandling:
    @patch("nix_builder_autoscaler.providers.haproxy.socket.socket")
    def test_unrecognized_slot_raises_haproxy_error(self, mock_socket_cls):
        mock_sock = MagicMock()
        mock_socket_cls.return_value = mock_sock
        mock_sock.recv.side_effect = [b"No such server.\n", b""]

        h = HAProxyRuntime("/tmp/test.sock", "all", "slot")
        with pytest.raises(HAProxyError, match="No such server"):
            h.set_slot_addr("slot999", "100.64.0.1")

    @patch("nix_builder_autoscaler.providers.haproxy.socket.socket")
    def test_socket_not_found_raises_haproxy_error(self, mock_socket_cls):
        mock_sock = MagicMock()
        mock_socket_cls.return_value = mock_sock
        mock_sock.connect.side_effect = FileNotFoundError("No such file")

        h = HAProxyRuntime("/tmp/nonexistent.sock", "all", "slot")
        with pytest.raises(HAProxyError, match="socket not found"):
            h.set_slot_addr("slot001", "100.64.0.1")

    @patch("nix_builder_autoscaler.providers.haproxy.socket.socket")
    def test_connection_refused_raises_haproxy_error(self, mock_socket_cls):
        mock_sock = MagicMock()
        mock_socket_cls.return_value = mock_sock
        mock_sock.connect.side_effect = ConnectionRefusedError("Connection refused")

        h = HAProxyRuntime("/tmp/test.sock", "all", "slot")
        with pytest.raises(HAProxyError, match="Connection refused"):
            h.enable_slot("slot001")

    @patch("nix_builder_autoscaler.providers.haproxy.socket.socket")
    def test_slot_session_count_missing_slot_raises(self, mock_socket_cls):
        mock_sock = MagicMock()
        mock_socket_cls.return_value = mock_sock
        mock_sock.recv.side_effect = [SHOW_STAT_CSV.encode(), b""]

        h = HAProxyRuntime("/tmp/test.sock", "all", "slot")
        with pytest.raises(HAProxyError, match="Slot not found"):
            h.slot_session_count("slot_nonexistent")