lmcache--lmcache/tests/tools/test_controller_zmq_benchmark.py

#!/usr/bin/env python3
"""
Tests for LMCache Controller ZMQ Benchmark Tool
"""

# SPDX-License-Identifier: Apache-2.0

# Standard
from typing import Dict
from unittest.mock import AsyncMock, MagicMock, patch

# Third Party
import pytest

# First Party
from lmcache.logging import init_logger
from lmcache.tools.controller_benchmark.benchmark import ZMQControllerBenchmark
from lmcache.tools.controller_benchmark.config import ZMQBenchmarkConfig
from lmcache.tools.controller_benchmark.handlers import OPERATION_HANDLERS
from lmcache.tools.controller_benchmark.handlers.admit import AdmitHandler
from lmcache.tools.controller_benchmark.handlers.base import (
    OperationHandler,
    SocketType,
)
from lmcache.tools.controller_benchmark.handlers.deregister import DeregisterHandler
from lmcache.tools.controller_benchmark.handlers.evict import EvictHandler
from lmcache.tools.controller_benchmark.handlers.heartbeat import HeartbeatHandler
from lmcache.tools.controller_benchmark.handlers.p2p_lookup import P2PLookupHandler
from lmcache.tools.controller_benchmark.handlers.register import RegisterHandler
from lmcache.v1.cache_controller.message import (
    BatchedKVOperationMsg,
    BatchedP2PLookupMsg,
    DeRegisterMsg,
    HeartbeatMsg,
    OpType,
    RegisterMsg,
)

logger = init_logger(__name__)


class TestZMQBenchmarkConfig:
    """Test ZMQBenchmarkConfig class"""

    def test_default_operations(self):
        """Test default operation distribution"""
        config = ZMQBenchmarkConfig(
            controller_pull_url="localhost:8100",
            controller_reply_url=None,
            duration=60,
            batch_size=50,
            num_instances=10,
            num_workers=5,
            num_locations=3,
            num_keys=1000,
        )

        assert "admit" in config.operations
        assert "evict" in config.operations
        assert "heartbeat" in config.operations
        assert abs(sum(config.operations.values()) - 100.0) < 0.01

    def test_custom_operations(self):
        """Test custom operation distribution"""
        config = ZMQBenchmarkConfig(
            controller_pull_url="localhost:8100",
            controller_reply_url="localhost:8101",
            duration=60,
            batch_size=50,
            num_instances=10,
            num_workers=5,
            num_locations=3,
            num_keys=1000,
            operations={"admit": 80.0, "evict": 20.0},
        )

        assert config.operations["admit"] == 80.0
        assert config.operations["evict"] == 20.0
        assert config.controller_pull_url == "localhost:8100"
        assert config.controller_reply_url == "localhost:8101"
        assert config.duration == 60
        assert config.batch_size == 50
        assert config.num_instances == 10
        assert config.num_workers == 5
        assert config.num_locations == 3
        assert config.num_keys == 1000

    def test_invalid_operation_percentage(self):
        """Test invalid operation percentage sum"""
        with pytest.raises(ValueError, match="must sum to 100"):
            ZMQBenchmarkConfig(
                controller_pull_url="localhost:8100",
                controller_reply_url=None,
                duration=60,
                batch_size=50,
                num_instances=10,
                num_workers=5,
                num_locations=3,
                num_keys=1000,
                operations={"admit": 50, "evict": 60},  # Sum > 100
            )


class TestOperationHandlers:
    """Test operation handlers (Strategy Pattern)"""

    @pytest.fixture
    def controller_zmq_benchmark(self):
        """Create a benchmark instance for testing"""
        config = ZMQBenchmarkConfig(
            controller_pull_url="localhost:8100",
            controller_reply_url="localhost:8101",
            duration=1,
            batch_size=10,
            num_instances=2,
            num_workers=2,
            num_locations=2,
            num_keys=10,
            num_hashes=5,
            register_first=False,
        )
        return ZMQControllerBenchmark(config)

    @pytest.fixture
    def test_data(self, controller_zmq_benchmark):
        """Generate test data"""
        return controller_zmq_benchmark.generate_test_data()

    def test_admit_handler(self, controller_zmq_benchmark, test_data):
        """Test AdmitHandler functionality"""
        handler = AdmitHandler()
        msg = handler.create_message(controller_zmq_benchmark, test_data)

        assert isinstance(msg, BatchedKVOperationMsg)
        assert len(msg.operations) == controller_zmq_benchmark.config.batch_size
        assert all(op.op_type == OpType.ADMIT for op in msg.operations)
        assert (
            handler.get_message_count(controller_zmq_benchmark)
            == controller_zmq_benchmark.config.batch_size
        )
        assert handler.socket_type == SocketType.PUSH

    def test_evict_handler(self, controller_zmq_benchmark, test_data):
        """Test EvictHandler functionality"""
        handler = EvictHandler()
        msg = handler.create_message(controller_zmq_benchmark, test_data)

        assert isinstance(msg, BatchedKVOperationMsg)
        assert len(msg.operations) == controller_zmq_benchmark.config.batch_size
        assert all(op.op_type == OpType.EVICT for op in msg.operations)
        assert (
            handler.get_message_count(controller_zmq_benchmark)
            == controller_zmq_benchmark.config.batch_size
        )
        assert handler.socket_type == SocketType.PUSH

    def test_heartbeat_handler(self, controller_zmq_benchmark, test_data):
        """Test HeartbeatHandler functionality"""
        handler = HeartbeatHandler()
        msg = handler.create_message(controller_zmq_benchmark, test_data)

        assert isinstance(msg, HeartbeatMsg)
        assert handler.get_message_count(controller_zmq_benchmark) == 1
        assert handler.socket_type == SocketType.HEARTBEAT

    def test_register_handler(self, controller_zmq_benchmark, test_data):
        """Test RegisterHandler functionality"""
        handler = RegisterHandler()
        msg = handler.create_message(controller_zmq_benchmark, test_data)

        assert isinstance(msg, RegisterMsg)
        assert handler.get_message_count(controller_zmq_benchmark) == 1
        assert handler.socket_type == SocketType.DEALER

    def test_deregister_handler(self, controller_zmq_benchmark, test_data):
        """Test DeregisterHandler functionality"""
        handler = DeregisterHandler()
        msg = handler.create_message(controller_zmq_benchmark, test_data)

        assert isinstance(msg, DeRegisterMsg)
        assert handler.get_message_count(controller_zmq_benchmark) == 1
        assert handler.socket_type == SocketType.PUSH

    def test_p2p_lookup_handler(self, controller_zmq_benchmark, test_data):
        """Test P2PLookupHandler functionality"""
        handler = P2PLookupHandler()
        msg = handler.create_message(controller_zmq_benchmark, test_data)

        assert isinstance(msg, BatchedP2PLookupMsg)
        assert len(msg.hashes) == controller_zmq_benchmark.config.num_hashes
        assert (
            handler.get_message_count(controller_zmq_benchmark)
            == controller_zmq_benchmark.config.num_hashes
        )
        assert handler.socket_type == SocketType.DEALER

    def test_operation_handler_registry(self):
        """Test operation handler registry"""
        assert "admit" in OPERATION_HANDLERS
        assert "evict" in OPERATION_HANDLERS
        assert "heartbeat" in OPERATION_HANDLERS
        assert "register" in OPERATION_HANDLERS
        assert "deregister" in OPERATION_HANDLERS
        assert "p2p_lookup" in OPERATION_HANDLERS

        for handler in OPERATION_HANDLERS.values():
            assert isinstance(handler, OperationHandler)


class TestZMQControllerBenchmark:
    """Test ZMQControllerBenchmark class"""

    @pytest.fixture
    def config(self):
        """Create test configuration"""
        return ZMQBenchmarkConfig(
            controller_pull_url="localhost:8100",
            controller_reply_url="localhost:8101",
            duration=1,
            batch_size=10,
            num_instances=2,
            num_workers=2,
            num_locations=2,
            num_keys=10,
            num_hashes=5,
            register_first=False,
        )

    @pytest.fixture
    def zmq_benchmark(self, config):
        """Create benchmark instance"""
        return ZMQControllerBenchmark(config)

    def test_generate_test_data(self, zmq_benchmark):
        """Test test data generation"""
        test_data = zmq_benchmark.generate_test_data()

        assert hasattr(test_data, "instances")
        assert hasattr(test_data, "workers")
        assert hasattr(test_data, "locations")
        assert hasattr(test_data, "keys")

        assert len(test_data.instances) == zmq_benchmark.config.num_instances
        assert len(test_data.workers) == zmq_benchmark.config.num_workers
        assert len(test_data.locations) == zmq_benchmark.config.num_locations
        assert len(test_data.keys) == zmq_benchmark.config.num_keys

    def test_sequence_number_increment(self, zmq_benchmark):
        """Test sequence number monotonic increment"""
        instance_id = "instance_0"
        worker_id = 0
        location = "location_0"

        seq1 = zmq_benchmark.get_next_sequence_number(instance_id, worker_id, location)
        seq2 = zmq_benchmark.get_next_sequence_number(instance_id, worker_id, location)
        seq3 = zmq_benchmark.get_next_sequence_number(instance_id, worker_id, location)

        assert seq1 == 0
        assert seq2 == 1
        assert seq3 == 2

    def test_sequence_number_isolation(self, zmq_benchmark):
        """Test sequence number isolation between different keys"""
        instance1, worker1, loc1 = "instance_0", 0, "location_0"
        instance2, worker2, loc2 = "instance_1", 1, "location_1"

        seq1a = zmq_benchmark.get_next_sequence_number(instance1, worker1, loc1)
        seq2a = zmq_benchmark.get_next_sequence_number(instance2, worker2, loc2)
        seq1b = zmq_benchmark.get_next_sequence_number(instance1, worker1, loc1)
        seq2b = zmq_benchmark.get_next_sequence_number(instance2, worker2, loc2)

        assert seq1a == 0
        assert seq2a == 0
        assert seq1b == 1
        assert seq2b == 1

    @pytest.mark.asyncio
    async def test_send_messages_success(self, zmq_benchmark):
        """Test successful message sending"""
        # Mock the socket
        mock_socket = AsyncMock()
        zmq_benchmark.push_socket = mock_socket

        test_data = zmq_benchmark.generate_test_data()
        handler = RegisterHandler()
        msg = handler.create_message(zmq_benchmark, test_data)

        latency = await zmq_benchmark.send_messages([msg])

        assert isinstance(latency, float)
        assert latency >= 0
        mock_socket.send_multipart.assert_called_once()

    @pytest.mark.asyncio
    async def test_send_messages_timeout(self, zmq_benchmark):
        """Test message sending timeout"""
        # Mock the socket to raise timeout
        mock_socket = AsyncMock()
        mock_socket.send_multipart.side_effect = Exception("Send timeout")
        zmq_benchmark.push_socket = mock_socket

        test_data = zmq_benchmark.generate_test_data()
        handler = RegisterHandler()
        msg = handler.create_message(zmq_benchmark, test_data)

        with pytest.raises(Exception, match="Send timeout"):
            await zmq_benchmark.send_messages([msg])

    @pytest.mark.asyncio
    async def test_send_request_success(self, zmq_benchmark):
        """Test successful request sending via DEALER socket"""
        # Mock the socket
        mock_socket = AsyncMock()
        mock_socket.send_multipart = AsyncMock()
        # DEALER receives: [empty_frame, payload]
        mock_socket.recv_multipart = AsyncMock(return_value=[b"", b"response"])
        zmq_benchmark.req_socket = mock_socket

        test_data = zmq_benchmark.generate_test_data()
        handler = P2PLookupHandler()
        msg = handler.create_message(zmq_benchmark, test_data)

        latency, response = await zmq_benchmark.send_request(msg)

        assert isinstance(latency, float)
        assert latency >= 0
        assert response == b"response"
        mock_socket.send_multipart.assert_called_once()
        mock_socket.recv_multipart.assert_called_once()

    @pytest.mark.asyncio
    async def test_send_request_timeout(self, zmq_benchmark):
        """Test request sending timeout via DEALER socket"""
        # Mock the socket to raise timeout
        mock_socket = AsyncMock()
        mock_socket.send_multipart = AsyncMock()
        mock_socket.recv_multipart = AsyncMock(side_effect=Exception("Request timeout"))
        zmq_benchmark.req_socket = mock_socket

        test_data = zmq_benchmark.generate_test_data()
        handler = P2PLookupHandler()
        msg = handler.create_message(zmq_benchmark, test_data)

        with pytest.raises(Exception, match="Request timeout"):
            await zmq_benchmark.send_request(msg)

    @pytest.mark.asyncio
    async def test_execute_operation_success(self, zmq_benchmark):
        """Test successful operation execution"""
        # Mock send_messages to avoid actual socket operations
        with patch.object(
            zmq_benchmark, "send_messages", AsyncMock(return_value=0.001)
        ):
            test_data = zmq_benchmark.generate_test_data()

            (
                msg_count,
                req_count,
                latency,
                error,
            ) = await zmq_benchmark._execute_operation("admit", test_data)

            assert msg_count == zmq_benchmark.config.batch_size
            assert req_count == 1
            assert latency > 0
            assert error is None

    @pytest.mark.asyncio
    async def test_execute_operation_unknown(self, zmq_benchmark):
        """Test unknown operation execution"""
        test_data = zmq_benchmark.generate_test_data()

        msg_count, req_count, latency, error = await zmq_benchmark._execute_operation(
            "unknown_op", test_data
        )

        assert msg_count == 0
        assert req_count == 0
        assert latency == 0.0
        assert isinstance(error, ValueError)

    def test_build_operation_distribution(self, zmq_benchmark):
        """Test operation distribution building"""
        operations = zmq_benchmark._build_operation_distribution()

        assert len(operations) > 0
        assert all(op in zmq_benchmark.config.operations for op in operations)

        # Verify distribution roughly matches percentages
        op_counts = {}
        for op in operations:
            op_counts[op] = op_counts.get(op, 0) + 1

        total_ops = len(operations)
        for op_name, percentage in zmq_benchmark.config.operations.items():
            expected_count = int(total_ops * percentage / 100)
            actual_count = op_counts.get(op_name, 0)
            # Allow some tolerance for random distribution
            assert abs(actual_count - expected_count) <= 2


class TestIntegration:
    """Integration tests for the ZMQ benchmark tool"""

    def test_operation_distribution(self):
        """Test operation distribution calculation"""
        config = ZMQBenchmarkConfig(
            controller_pull_url="localhost:8100",
            controller_reply_url=None,
            duration=60,
            batch_size=50,
            num_instances=10,
            num_workers=5,
            num_locations=3,
            num_keys=1000,
            operations={"admit": 50, "evict": 30, "heartbeat": 20},
        )

        zmq_benchmark = ZMQControllerBenchmark(config)
        # Verify benchmark was created with correct config
        assert zmq_benchmark.config.batch_size == 50

        # Build operation distribution
        operations = []
        total_ops = 1000
        for op_name, percentage in config.operations.items():
            count = int(total_ops * percentage / 100)
            operations.extend([op_name] * count)

        # Verify distribution
        assert len(operations) == total_ops

        op_counts: Dict[str, int] = {}
        for op in operations:
            op_counts[op] = op_counts.get(op, 0) + 1

        assert op_counts["admit"] == 500
        assert op_counts["evict"] == 300
        assert op_counts["heartbeat"] == 200

    @pytest.mark.asyncio
    async def test_benchmark_setup_and_cleanup(self):
        """Test benchmark setup and cleanup"""
        config = ZMQBenchmarkConfig(
            controller_pull_url="localhost:8100",
            controller_reply_url="localhost:8101",
            duration=1,
            batch_size=10,
            num_instances=2,
            num_workers=2,
            num_locations=2,
            num_keys=10,
            register_first=False,
        )

        zmq_benchmark = ZMQControllerBenchmark(config)

        # Mock ZMQ context and sockets
        with patch(
            "lmcache.tools.controller_benchmark.benchmark.get_zmq_context"
        ) as mock_context:
            mock_ctx = MagicMock()
            mock_context.return_value = mock_ctx

            mock_socket = MagicMock()
            mock_ctx.socket.return_value = mock_socket

            await zmq_benchmark.setup()

            # Verify setup was called
            mock_context.assert_called_once()

            # Test cleanup
            zmq_benchmark.cleanup()
            # Verify sockets were closed
            assert mock_socket.close.called

    def test_results_structure(self):
        """Test results structure initialization"""
        config = ZMQBenchmarkConfig(
            controller_pull_url="localhost:8100",
            controller_reply_url="localhost:8101",
            duration=1,
            batch_size=10,
            num_instances=2,
            num_workers=2,
            num_locations=2,
            num_keys=10,
            register_first=False,
        )

        zmq_benchmark = ZMQControllerBenchmark(config)
        results = zmq_benchmark.results

        assert hasattr(results, "operations")
        assert hasattr(results, "memory_usage")
        assert hasattr(results, "total_messages")
        assert hasattr(results, "total_requests")
        assert hasattr(results, "overall_qps")
        assert hasattr(results, "overall_rps")

        assert isinstance(results.operations, dict)
        assert isinstance(results.memory_usage, list)
        assert results.total_messages == 0
        assert results.total_requests == 0

    def test_normalize_heartbeat_url_localhost_connection(self):
        """Test heartbeat URL normalization when connecting to localhost"""
        config = ZMQBenchmarkConfig(
            controller_pull_url="127.0.0.1:8100",
            controller_reply_url="127.0.0.1:8101",
            duration=1,
            batch_size=10,
            num_instances=2,
            num_workers=2,
            num_locations=2,
            num_keys=10,
            register_first=False,
        )

        zmq_benchmark = ZMQControllerBenchmark(config)

        # When connecting locally, remote IPs should be replaced with 127.0.0.1
        assert (
            zmq_benchmark._normalize_heartbeat_url("10.0.0.5:7557") == "127.0.0.1:7557"
        )
        assert (
            zmq_benchmark._normalize_heartbeat_url("192.168.1.100:7557")
            == "127.0.0.1:7557"
        )
        # localhost should remain unchanged
        assert (
            zmq_benchmark._normalize_heartbeat_url("127.0.0.1:7557") == "127.0.0.1:7557"
        )
        assert (
            zmq_benchmark._normalize_heartbeat_url("localhost:7557") == "localhost:7557"
        )

    def test_normalize_heartbeat_url_remote_connection(self):
        """Test heartbeat URL normalization when connecting to remote host"""
        config = ZMQBenchmarkConfig(
            controller_pull_url="10.0.0.5:8100",
            controller_reply_url="10.0.0.5:8101",
            duration=1,
            batch_size=10,
            num_instances=2,
            num_workers=2,
            num_locations=2,
            num_keys=10,
            register_first=False,
        )

        zmq_benchmark = ZMQControllerBenchmark(config)

        # When connecting to remote, should keep the original IP
        assert (
            zmq_benchmark._normalize_heartbeat_url("10.0.0.5:7557") == "10.0.0.5:7557"
        )
        assert (
            zmq_benchmark._normalize_heartbeat_url("192.168.1.100:7557")
            == "192.168.1.100:7557"
        )