wehub-resource-sync
124 lines
4.0 KiB
Python
124 lines
4.0 KiB
Python
# SPDX-License-Identifier: Apache-2.0
|
|
"""
|
|
Tests for _run_with_disconnect_guard in server module.
|
|
|
|
Tests cover:
|
|
- Normal completion returns result
|
|
- Client disconnect cancels task
|
|
- Fast completion has no overhead from polling
|
|
"""
|
|
|
|
import asyncio
|
|
from unittest.mock import AsyncMock
|
|
|
|
import pytest
|
|
|
|
|
|
class TestDisconnectGuard:
|
|
"""Tests for _run_with_disconnect_guard."""
|
|
|
|
@pytest.fixture
|
|
def mock_request_connected(self):
|
|
"""Mock HTTP request that stays connected."""
|
|
request = AsyncMock()
|
|
request.is_disconnected = AsyncMock(return_value=False)
|
|
return request
|
|
|
|
@pytest.fixture
|
|
def mock_request_disconnects(self):
|
|
"""Mock HTTP request that disconnects after first check."""
|
|
request = AsyncMock()
|
|
request.is_disconnected = AsyncMock(side_effect=[False, True])
|
|
return request
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_normal_completion(self, mock_request_connected):
|
|
"""Test that normal completion returns result."""
|
|
from omlx.server import _run_with_disconnect_guard
|
|
|
|
async def fake_generate():
|
|
return "result"
|
|
|
|
result = await _run_with_disconnect_guard(
|
|
mock_request_connected, fake_generate(), poll_interval=0.1
|
|
)
|
|
assert result == "result"
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_disconnect_cancels_task(self, mock_request_disconnects):
|
|
"""Test that disconnect cancels the running task."""
|
|
from omlx.server import _run_with_disconnect_guard
|
|
|
|
cancel_detected = False
|
|
|
|
async def slow_generate():
|
|
nonlocal cancel_detected
|
|
try:
|
|
await asyncio.sleep(10)
|
|
return "should not reach"
|
|
except asyncio.CancelledError:
|
|
cancel_detected = True
|
|
raise
|
|
|
|
result = await _run_with_disconnect_guard(
|
|
mock_request_disconnects, slow_generate(), poll_interval=0.1
|
|
)
|
|
|
|
assert result is None # Client disconnected
|
|
assert cancel_detected # Task was actually cancelled
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_fast_completion_no_disconnect_check(self, mock_request_connected):
|
|
"""Test that fast completions finish without disconnect check."""
|
|
from omlx.server import _run_with_disconnect_guard
|
|
|
|
async def fast_generate():
|
|
return "fast_result"
|
|
|
|
result = await _run_with_disconnect_guard(
|
|
mock_request_connected, fast_generate(), poll_interval=1.0
|
|
)
|
|
assert result == "fast_result"
|
|
# Task completed before poll interval, so is_disconnected should not be called
|
|
mock_request_connected.is_disconnected.assert_not_called()
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_disconnect_during_long_generation(self):
|
|
"""Test disconnect detection during a long-running generation."""
|
|
from omlx.server import _run_with_disconnect_guard
|
|
|
|
call_count = 0
|
|
|
|
async def delayed_disconnect():
|
|
nonlocal call_count
|
|
call_count += 1
|
|
# Stay connected for 2 checks, then disconnect
|
|
return call_count > 2
|
|
|
|
mock_request = AsyncMock()
|
|
mock_request.is_disconnected = delayed_disconnect
|
|
|
|
async def slow_generate():
|
|
await asyncio.sleep(10)
|
|
return "should not reach"
|
|
|
|
result = await _run_with_disconnect_guard(
|
|
mock_request, slow_generate(), poll_interval=0.05
|
|
)
|
|
|
|
assert result is None
|
|
assert call_count == 3 # Connected, connected, disconnected
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_task_exception_propagates(self, mock_request_connected):
|
|
"""Test that task exceptions propagate correctly."""
|
|
from omlx.server import _run_with_disconnect_guard
|
|
|
|
async def failing_generate():
|
|
raise ValueError("generation failed")
|
|
|
|
with pytest.raises(ValueError, match="generation failed"):
|
|
await _run_with_disconnect_guard(
|
|
mock_request_connected, failing_generate(), poll_interval=0.1
|
|
)
|