1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
|
"""
Copyright (c) 2023 Proton AG
This file is part of Proton.
Proton is free software: you can redistribute it and/or modify
it under the terms of the GNU General Public License as published by
the Free Software Foundation, either version 3 of the License, or
(at your option) any later version.
Proton is distributed in the hope that it will be useful,
but WITHOUT ANY WARRANTY; without even the implied warranty of
MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
GNU General Public License for more details.
You should have received a copy of the GNU General Public License
along with ProtonVPN. If not, see <https://www.gnu.org/licenses/>.
"""
from unittest.mock import Mock, AsyncMock
import pytest
import asyncio
import os
import time
import unittest
from proton.session import Session
from proton.session.transports.auto import AutoTransport
from proton.session.transports.requests import RequestsTransport
from proton.session.exceptions import ProtonAPINotReachable
class TestAuto(unittest.IsolatedAsyncioTestCase):
def setUp(self):
self._env_backup = os.environ.copy()
def tearDown(self):
os.environ = self._env_backup
async def test_auto_works_on_prod(self):
os.environ['PROTON_API_ENVIRONMENT'] = 'prod'
s = Session()
s.transport_factory = AutoTransport
assert await s.async_api_request('/tests/ping') == {'Code': 1000}
async def test_auto_transport_is_not_available_when_all_transports_choices_time_out_pinging_rest_api(self):
mock_transport_type = Mock()
transport_timeout = 0.001
auto_transport = AutoTransport(
session=Session(),
transport_choices=[(0, mock_transport_type)],
transport_timeout=transport_timeout
)
mock_transport = Mock()
mock_transport_type.return_value = mock_transport
# Force a timeout from `/tests/ping` when checking if the transport is available.
async def force_transport_timeout(url):
await asyncio.sleep(transport_timeout + 1)
mock_transport.async_api_request.side_effect = force_transport_timeout
with pytest.raises(ProtonAPINotReachable):
await auto_transport.find_available_transport()
mock_transport.async_api_request.assert_called_once_with('/tests/ping')
assert not auto_transport.is_available
async def test_auto_transport_is_not_available_when_all_transport_choices_receive_an_unexpected_ping_response(self):
mock_transport_type = Mock()
auto_transport = AutoTransport(
session=Session(),
transport_choices=[(0, mock_transport_type)]
)
mock_transport = Mock()
mock_transport_type.return_value = mock_transport
# Force an unexpected response from `/tests/ping` when checking if the transport is available.
async def force_unexpected_ping_response(url):
return "foobar"
mock_transport.async_api_request.side_effect = force_unexpected_ping_response
with pytest.raises(ProtonAPINotReachable):
await auto_transport.find_available_transport()
mock_transport.async_api_request.assert_called_once_with('/tests/ping')
assert not auto_transport.is_available
|
