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
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
|
"""
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/>.
"""
import unittest
from io import StringIO
from unittest.mock import patch, AsyncMock, Mock
from proton.session import Session
from proton.session.formdata import FormData, FormField
from proton.session.transports import AiohttpTransport
from proton.session.transports.aiohttp import FormDataTransformer
from proton.session.transports.base import RawResponse
HTTP_STATUS_OK = 200
HTTP_STATUS_NOT_MODIFIED = 304
CODE_SUCCESS = 1000
class TestAiohttpTransport(unittest.IsolatedAsyncioTestCase):
@patch("proton.session.transports.aiohttp.aiohttp.ClientSession.post")
async def test_async_api_request_posts_form_data_with_data_param(self, post_mock):
session = Session()
form_data_transformer_mock = Mock(spec=FormDataTransformer)
aiohttp_transport = AiohttpTransport(session, form_data_transformer_mock)
# Mock POST response.
post_mock.return_value.__aenter__.return_value.status = HTTP_STATUS_OK
post_mock.return_value.__aenter__.return_value.headers = {"content-type": "application/json"}
post_mock.return_value.__aenter__.return_value.json = AsyncMock(
return_value={"Code": CODE_SUCCESS}
)
# Form data to be posted.
form_data = FormData()
form_data.add(FormField(name="foo", value="bar"))
# SUT.
await aiohttp_transport.async_api_request("/endpoint", data=form_data)
# Assert that the form data has been transformed to aiohttp.FormData.
form_data_transformer_mock.to_aiohttp_form_data.assert_called_once_with(form_data)
expected_payload_to_be_posted = form_data_transformer_mock.to_aiohttp_form_data.return_value
# Assert that the POST call is done with the transformed form data.
post_mock.assert_called_once()
posted_payload = post_mock.call_args.kwargs["data"]
assert posted_payload is expected_payload_to_be_posted
class TestFormDataTransformer(unittest.TestCase):
@patch("proton.session.transports.aiohttp.aiohttp.FormData")
def test_to_aiohttp_form_data(self, _aiohttp_form_data_mock):
form_data_transformer = FormDataTransformer()
# Form data to be transformed:
form_data = FormData()
# Add a simple field to the form.
first_field_name, first_field_value = "foo", "bar"
form_data.add(FormField(name=first_field_name, value=first_field_value))
# Add a file to the form.
second_field_name = "file"
second_field_value = StringIO("File content.")
second_field_filename = "file.txt"
second_field_content_type = "text/plain"
form_data.add(FormField(
name=second_field_name, value=second_field_value,
filename=second_field_filename, content_type=second_field_content_type
))
# SUT.
result = form_data_transformer.to_aiohttp_form_data(form_data)
# Assert that aiohttp.FormData was created with the form data passed above.
assert result.add_field.call_count == 2
assert result.add_field.call_args_list[0].kwargs == {
"name": first_field_name, "value": first_field_value,
"content_type": None, "filename": None
}
assert result.add_field.call_args_list[1].kwargs == {
"name": second_field_name, "value": second_field_value,
"content_type": second_field_content_type, "filename": second_field_filename
}
class TestAiohttpTransportRawResult(unittest.IsolatedAsyncioTestCase):
def _setup(self, get_mock, status, headers, json):
# Mock the GET response
get_mock.return_value.__aenter__.return_value.status = status
get_mock.return_value.__aenter__.return_value.headers = headers
get_mock.return_value.__aenter__.return_value.json = AsyncMock(
return_value=json
)
session = Session()
aiohttp_transport = AiohttpTransport(session)
return session, aiohttp_transport
@patch("proton.session.transports.aiohttp.aiohttp.ClientSession.get")
async def test_async_api_request_get_raw(self, get_mock):
# Setup
session, aiohttp_transport = self._setup(get_mock,
status=HTTP_STATUS_OK,
headers={"content-type": "application/json"},
json={"Code": CODE_SUCCESS})
# Test
response = await aiohttp_transport.async_api_request("/endpoint", return_raw=True)
# Checks
assert isinstance(response, RawResponse), "The response should be a RawResponse object."
assert response.status_code == HTTP_STATUS_OK
assert response.find_first_header("content-type") == "application/json"
assert response.json == {"Code": CODE_SUCCESS}
get_mock.assert_called_once()
@patch("proton.session.transports.aiohttp.aiohttp.ClientSession.get")
async def test_async_api_request_last_modified(self, get_mock):
# Setup
session, aiohttp_transport = self._setup(get_mock,
status=HTTP_STATUS_NOT_MODIFIED,
headers={},
json=None)
# Test
response = await aiohttp_transport.async_api_request("/endpoint", return_raw=True)
# Checks
assert isinstance(response, RawResponse), "The response should be a RawResponse object."
assert response.status_code == HTTP_STATUS_NOT_MODIFIED
assert response.find_first_header("content-type", None) is None
assert response.json is None
get_mock.assert_called_once()
|
