|
| 1 | +from unittest.mock import AsyncMock, MagicMock, Mock |
| 2 | + |
| 3 | +import pytest |
| 4 | +from aleph_message.models import ItemHash |
| 5 | + |
| 6 | +from aleph.vm.orchestrator.status import check_internet |
| 7 | + |
| 8 | + |
| 9 | +@pytest.mark.asyncio |
| 10 | +async def test_check_internet_no_server_header(): |
| 11 | + vm_id = ItemHash("cafecafecafecafecafecafecafecafecafecafecafecafecafecafecafecafe") |
| 12 | + |
| 13 | + mock_session = Mock() |
| 14 | + mock_session.get = MagicMock() |
| 15 | + mock_session.get.__aenter__.return_value.json = AsyncMock(return_value={"result": 200}) |
| 16 | + |
| 17 | + # A "Server" header is always expected in the result from the VM. |
| 18 | + # If it is not present, the diagnostic VM is not working correctly |
| 19 | + # and an error must be raised. |
| 20 | + with pytest.raises(ValueError): |
| 21 | + await check_internet(mock_session, vm_id) |
| 22 | + |
| 23 | + |
| 24 | +@pytest.mark.asyncio |
| 25 | +async def test_check_internet_wrong_result_code(): |
| 26 | + vm_id = ItemHash("cafecafecafecafecafecafecafecafecafecafecafecafecafecafecafecafe") |
| 27 | + |
| 28 | + mock_session = Mock() |
| 29 | + mock_session.get = MagicMock() |
| 30 | + |
| 31 | + mock_session.get.return_value.__aenter__.return_value.json = AsyncMock( |
| 32 | + return_value={"result": 200, "headers": {"Server": "nginx"}} |
| 33 | + ) |
| 34 | + assert await check_internet(mock_session, vm_id) is True |
| 35 | + |
| 36 | + mock_session.get.return_value.__aenter__.return_value.json = AsyncMock( |
| 37 | + return_value={"result": 400, "headers": {"Server": "nginx"}} |
| 38 | + ) |
| 39 | + assert await check_internet(mock_session, vm_id) is False |
0 commit comments