diff --git a/aiosonic/exceptions.py b/aiosonic/exceptions.py index c43f29b..d2a41d8 100644 --- a/aiosonic/exceptions.py +++ b/aiosonic/exceptions.py @@ -34,3 +34,8 @@ class HttpParsingError(Exception): # Redirects class MaxRedirects(Exception): pass + + +# HTTP2 +class MissingEvent(Exception): + pass diff --git a/aiosonic/http2.py b/aiosonic/http2.py index 5527c07..3d3ed5e 100644 --- a/aiosonic/http2.py +++ b/aiosonic/http2.py @@ -5,6 +5,7 @@ import h2.events +from aiosonic.exceptions import MissingEvent from aiosonic.types import ParamsType from aiosonic.types import ParsedBodyType @@ -79,44 +80,49 @@ async def request(self, headers: ParamsType, body: Optional[ParsedBodyType]): async def reader_t(self): """Reader task.""" read_size = 16000 - h2conn = self.h2conn while True: data = await asyncio.wait_for(self.reader.read(read_size), 3) - events = h2conn.receive_data(data) + events = self.h2conn.receive_data(data) if not events: continue else: - for event in events: - if isinstance(event, h2.events.StreamEnded): - self.requests[event.stream_id]['future'].set_result( - self.requests[event.stream_id]['body']) - elif isinstance(event, h2.events.DataReceived): - self.requests[event.stream_id]['body'] += event.data - - if (event.stream_id in h2conn.streams - and not h2conn.streams[event.stream_id].closed): - h2conn.increment_flow_control_window( - event.flow_controlled_length, event.stream_id) - h2conn.increment_flow_control_window( - event.flow_controlled_length) - elif isinstance(event, h2.events.ResponseReceived): - self.requests[ - event.stream_id]['headers'] = event.headers - elif isinstance(event, ( - h2.events.WindowUpdated, - h2.events.PingReceived, - h2.events.RemoteSettingsChanged, - h2.events.SettingsAcknowledged - )): - pass - else: - raise Exception( - f'another event {event.__class__.__name__}') + self.handle_events(events) await self.writer_q.put(True) + def handle_events(self, events): + """Handle http2 events.""" + h2conn = self.h2conn + + for event in events: + if isinstance(event, h2.events.StreamEnded): + self.requests[event.stream_id]['future'].set_result( + self.requests[event.stream_id]['body']) + elif isinstance(event, h2.events.DataReceived): + self.requests[event.stream_id]['body'] += event.data + + if (event.stream_id in h2conn.streams + and not h2conn.streams[event.stream_id].closed): + h2conn.increment_flow_control_window( + event.flow_controlled_length, event.stream_id) + h2conn.increment_flow_control_window( + event.flow_controlled_length) + elif isinstance(event, h2.events.ResponseReceived): + self.requests[ + event.stream_id]['headers'] = event.headers + elif isinstance(event, ( + h2.events.WindowUpdated, + h2.events.PingReceived, + h2.events.RemoteSettingsChanged, + h2.events.SettingsAcknowledged + )): + pass + else: + raise MissingEvent( + f'another event {event.__class__.__name__}') + async def writer_t(self): """Writer task.""" h2conn = self.h2conn diff --git a/tests/test_aiosonic.py b/tests/test_aiosonic.py index db810b6..4cfdab4 100644 --- a/tests/test_aiosonic.py +++ b/tests/test_aiosonic.py @@ -5,19 +5,21 @@ import pytest import aiosonic +from aiosonic import _get_url_parsed +from aiosonic import HttpResponse +from aiosonic.connectors import TCPConnector +from aiosonic.connection import Connection from aiosonic.exceptions import ConnectTimeout from aiosonic.exceptions import ReadTimeout from aiosonic.exceptions import RequestTimeout from aiosonic.exceptions import MaxRedirects from aiosonic.exceptions import HttpParsingError from aiosonic.exceptions import MissingWriterException +from aiosonic.exceptions import MissingEvent from aiosonic.exceptions import ConnectionPoolAcquireTimeout -from aiosonic.connectors import TCPConnector -from aiosonic.connection import Connection +from aiosonic.http2 import Http2Handler from aiosonic.pools import CyclicQueuePool from aiosonic.timeout import Timeouts -from aiosonic import _get_url_parsed -from aiosonic import HttpResponse @pytest.mark.asyncio @@ -618,3 +620,22 @@ async def test_json_response_parsing(): response._set_header('content-type', 'application/json; charset=utf-8') response.body = b'{"foo": "bar"}' assert (await response.json()) == {'foo': 'bar'} + + +class WrongEvent: + pass + + +@pytest.mark.asyncio +async def test_http2_wrong_event(mocker): + """Test json response parsing.""" + mocker.patch('aiosonic.http2.Http2Handler.__init__', lambda x: None) + mocker.patch('aiosonic.http2.Http2Handler.h2conn') + + handler = Http2Handler() + + async def coro(): + pass + + with pytest.raises(MissingEvent): + await handler.handle_events([WrongEvent])