modelexpress-client 0.5.1

# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
# SPDX-License-Identifier: Apache-2.0

"""Tests for source publication and heartbeat signaling."""

import threading
import time
from unittest.mock import MagicMock, call, patch

import pytest

from modelexpress.metadata.publisher import PublisherThread


@pytest.fixture
def mx_client():
    client = MagicMock()
    client.update_status.return_value = True
    return client


@pytest.fixture
def nixl_manager():
    manager = MagicMock()
    manager.is_healthy.return_value = True
    return manager


@pytest.fixture
def heartbeat(mx_client, nixl_manager):
    with patch.dict("os.environ", {"MX_HEARTBEAT_INTERVAL_SECS": "1"}):
        hb = PublisherThread(
            mx_client=mx_client,
            mx_source_id="abc123",
            worker_id="w1",
            worker_rank=0,
            nixl_manager=nixl_manager,
        )
    yield hb
    hb.stop()


class TestHeartbeatSendsReady:
    def test_sends_ready_when_healthy(self, heartbeat, mx_client, nixl_manager):
        heartbeat.start()
        time.sleep(1.5)
        heartbeat.stop()

        calls = mx_client.update_status.call_args_list
        assert len(calls) >= 1
        assert calls[0] == call(
            mx_source_id="abc123",
            worker_id="w1",
            worker_rank=0,
            status=2,  # SOURCE_STATUS_READY
        )

    def test_skips_when_unhealthy(self, heartbeat, mx_client, nixl_manager):
        nixl_manager.is_healthy.return_value = False
        heartbeat.start()
        time.sleep(1.5)
        heartbeat.stop()

        # Only the _mark_stale call from stop(), no READY calls
        ready_calls = [
            c for c in mx_client.update_status.call_args_list
            if c == call(
                mx_source_id="abc123",
                worker_id="w1",
                worker_rank=0,
                status=2,
            )
        ]
        assert len(ready_calls) == 0

    def test_multiple_ticks_refresh_updated_at(self, heartbeat, mx_client):
        heartbeat.start()
        time.sleep(2.5)
        heartbeat.stop()

        # At 1s interval, 2.5s sleep should give at least 2 READY calls
        ready_calls = [
            c for c in mx_client.update_status.call_args_list
            if c == call(
                mx_source_id="abc123",
                worker_id="w1",
                worker_rank=0,
                status=2,
            )
        ]
        assert len(ready_calls) >= 2


class TestPublisherPublishAndReady:
    def test_first_tick_publishes_then_sends_ready(self, mx_client, nixl_manager):
        publisher = PublisherThread(
            mx_client=mx_client,
            worker_id="w1",
            worker_rank=0,
            nixl_manager=nixl_manager,
            publish_fn=lambda: "abc123",
            interval_secs=1,
        )

        publisher._tick()

        assert publisher.mx_source_id == "abc123"
        assert mx_client.update_status.call_args_list == [
            call(
                mx_source_id="abc123",
                worker_id="w1",
                worker_rank=0,
                status=2,
            )
        ]

    def test_ready_gate_blocks_publish(self, mx_client, nixl_manager):
        publish = MagicMock(return_value="abc123")
        publisher = PublisherThread(
            mx_client=mx_client,
            worker_id="w1",
            worker_rank=0,
            nixl_manager=nixl_manager,
            publish_fn=publish,
            ready_fn=lambda: False,
            interval_secs=1,
        )

        publisher._tick()

        publish.assert_not_called()
        mx_client.update_status.assert_not_called()

    def test_ready_gate_is_polled_before_heartbeat_interval(self, mx_client, nixl_manager):
        publisher = PublisherThread(
            mx_client=mx_client,
            worker_id="w1",
            worker_rank=0,
            nixl_manager=nixl_manager,
            publish_fn=MagicMock(return_value="abc123"),
            ready_fn=lambda: False,
            interval_secs=30,
        )
        waits = []

        def stop_after_wait(timeout):
            waits.append(timeout)
            publisher._stop_event.set()

        with patch.object(publisher._stop_event, "wait", side_effect=stop_after_wait):
            publisher._run()

        assert waits == [5]

    def test_can_stop_after_publish_without_heartbeat(self, mx_client, nixl_manager):
        ready = MagicMock(return_value=True)
        publish = MagicMock(return_value="abc123")
        publisher = PublisherThread(
            mx_client=mx_client,
            worker_id="w1",
            worker_rank=0,
            nixl_manager=nixl_manager,
            publish_fn=publish,
            ready_fn=ready,
            heartbeat_after_publish=False,
            interval_secs=1,
        )

        publisher._tick()
        publisher._tick()

        assert publisher.mx_source_id == "abc123"
        mx_client.update_status.assert_not_called()
        assert publisher._stop_event.is_set()
        ready.assert_called_once_with()
        publish.assert_called_once_with()


class TestHeartbeatStop:
    def test_stop_marks_stale(self, heartbeat, mx_client):
        heartbeat.start()
        time.sleep(1.5)  # Let at least one READY go through
        heartbeat.stop()

        stale_calls = [
            c for c in mx_client.update_status.call_args_list
            if c == call(
                mx_source_id="abc123",
                worker_id="w1",
                worker_rank=0,
                status=3,  # SOURCE_STATUS_STALE
            )
        ]
        assert len(stale_calls) == 1

    def test_stop_without_ready_skips_stale(self, heartbeat, mx_client, nixl_manager):
        nixl_manager.is_healthy.return_value = False
        heartbeat.start()
        time.sleep(1.5)
        heartbeat.stop()

        # Never became READY, so _mark_stale should not send STALE
        stale_calls = [
            c for c in mx_client.update_status.call_args_list
            if c == call(
                mx_source_id="abc123",
                worker_id="w1",
                worker_rank=0,
                status=3,
            )
        ]
        assert len(stale_calls) == 0

    def test_stop_is_idempotent(self, heartbeat, mx_client):
        heartbeat.start()
        time.sleep(1.5)
        heartbeat.stop()
        heartbeat.stop()  # Second stop should not send another STALE

        stale_calls = [
            c for c in mx_client.update_status.call_args_list
            if c == call(
                mx_source_id="abc123",
                worker_id="w1",
                worker_rank=0,
                status=3,
            )
        ]
        assert len(stale_calls) == 1


class TestHeartbeatOnExit:
    def test_on_exit_marks_stale(self, heartbeat, mx_client):
        heartbeat.start()
        time.sleep(1.5)
        heartbeat._on_exit()

        stale_calls = [
            c for c in mx_client.update_status.call_args_list
            if c == call(
                mx_source_id="abc123",
                worker_id="w1",
                worker_rank=0,
                status=3,
            )
        ]
        assert len(stale_calls) == 1

    def test_on_exit_swallows_errors(self, heartbeat, mx_client):
        heartbeat.start()
        time.sleep(1.5)
        mx_client.update_status.side_effect = RuntimeError("connection lost")
        heartbeat._on_exit()  # Should not raise

    def test_on_exit_orders_stale_after_inflight_ready(self, heartbeat, mx_client):
        ready_started = threading.Event()
        release_ready = threading.Event()

        def update_status(**kwargs):
            if kwargs["status"] == 2:
                ready_started.set()
                release_ready.wait(timeout=5)
            return True

        mx_client.update_status.side_effect = update_status
        heartbeat.start()
        assert ready_started.wait(timeout=5)

        exit_thread = threading.Thread(target=heartbeat._on_exit)
        exit_thread.start()
        assert heartbeat._stop_event.wait(timeout=5)
        release_ready.set()
        exit_thread.join(timeout=5)

        assert not exit_thread.is_alive()
        statuses = [call.kwargs["status"] for call in mx_client.update_status.call_args_list]
        assert statuses == [2, 3]


class TestHeartbeatDaemon:
    def test_thread_is_daemon(self, heartbeat):
        heartbeat.start()
        assert heartbeat._thread.daemon is True

    def test_update_status_error_does_not_crash_thread(self, heartbeat, mx_client):
        mx_client.update_status.side_effect = [
            RuntimeError("transient"),  # First tick fails
            True,                        # Second tick succeeds
            True,
        ]
        heartbeat.start()
        time.sleep(2.5)

        assert heartbeat._thread.is_alive()