lightrag-cn-git-d5efd3/tests/kg/test_unified_lock_safety.py

"""
Tests for UnifiedLock safety when lock is None.

This test module verifies that get_internal_lock() and get_data_init_lock()
raise RuntimeError when shared data is not initialized, preventing false
security and potential race conditions.

Design: The None check has been moved from UnifiedLock.__aenter__/__enter__
to the lock factory functions (get_internal_lock, get_data_init_lock) for
early failure detection.

Critical Bug 1 (Fixed): When self._lock is None, the code would fail with
AttributeError. Now the check is in factory functions for clearer errors.

Critical Bug 2: In __aexit__, when async_lock.release() fails, the error
recovery logic would attempt to release it again, causing double-release issues.
"""

from unittest.mock import MagicMock, AsyncMock

import pytest

from lightrag.kg.shared_storage import (
    UnifiedLock,
    get_internal_lock,
    get_data_init_lock,
    finalize_share_data,
)


class TestUnifiedLockSafety:
    """Test suite for UnifiedLock None safety checks."""

    def setup_method(self):
        """Ensure shared data is finalized before each test."""
        finalize_share_data()

    def teardown_method(self):
        """Clean up after each test."""
        finalize_share_data()

    def test_get_internal_lock_raises_when_not_initialized(self):
        """
        Test that get_internal_lock() raises RuntimeError when shared data is not initialized.

        Scenario: Call get_internal_lock() before initialize_share_data() is called.
        Expected: RuntimeError raised with clear error message.

        This test verifies the None check has been moved to the factory function.
        """
        with pytest.raises(
            RuntimeError, match="Shared data not initialized.*initialize_share_data"
        ):
            get_internal_lock()

    def test_get_data_init_lock_raises_when_not_initialized(self):
        """
        Test that get_data_init_lock() raises RuntimeError when shared data is not initialized.

        Scenario: Call get_data_init_lock() before initialize_share_data() is called.
        Expected: RuntimeError raised with clear error message.

        This test verifies the None check has been moved to the factory function.
        """
        with pytest.raises(
            RuntimeError, match="Shared data not initialized.*initialize_share_data"
        ):
            get_data_init_lock()

    @pytest.mark.offline
    async def test_aexit_no_double_release_on_async_lock_failure(self):
        """
        Test that __aexit__ doesn't attempt to release async_lock twice when it fails.

        Scenario: async_lock.release() fails during normal release.
        Expected: Recovery logic should NOT attempt to release async_lock again,
                  preventing double-release issues.

        This tests Bug 2 fix: async_lock_released tracking prevents double release.
        """
        # Create mock locks
        main_lock = MagicMock()
        main_lock.acquire = MagicMock()
        main_lock.release = MagicMock()

        async_lock = AsyncMock()
        async_lock.acquire = AsyncMock()

        # Make async_lock.release() fail
        release_call_count = 0

        def mock_release_fail():
            nonlocal release_call_count
            release_call_count += 1
            raise RuntimeError("Async lock release failed")

        async_lock.release = MagicMock(side_effect=mock_release_fail)

        # Create UnifiedLock with both locks (sync mode with async_lock)
        lock = UnifiedLock(
            lock=main_lock,
            is_async=False,
            name="test_double_release",
            enable_logging=False,
        )
        lock._async_lock = async_lock

        # Try to use the lock - should fail during __aexit__
        try:
            async with lock:
                pass
        except RuntimeError as e:
            # Should get the async lock release error
            assert "Async lock release failed" in str(e)

        # Verify async_lock.release() was called only ONCE, not twice
        assert (
            release_call_count == 1
        ), f"async_lock.release() should be called only once, but was called {release_call_count} times"

        # Main lock should have been released successfully
        main_lock.release.assert_called_once()