helpers.py (4865B)
1 # -*- coding: utf-8 -*- 2 """ 3 helpers 4 ~~~~~~~ 5 6 This module contains helpers for the h2 tests. 7 """ 8 from hyperframe.frame import ( 9 HeadersFrame, DataFrame, SettingsFrame, WindowUpdateFrame, PingFrame, 10 GoAwayFrame, RstStreamFrame, PushPromiseFrame, PriorityFrame, 11 ContinuationFrame, AltSvcFrame 12 ) 13 from hpack.hpack import Encoder 14 15 16 SAMPLE_SETTINGS = { 17 SettingsFrame.HEADER_TABLE_SIZE: 4096, 18 SettingsFrame.ENABLE_PUSH: 1, 19 SettingsFrame.MAX_CONCURRENT_STREAMS: 2, 20 } 21 22 23 class FrameFactory(object): 24 """ 25 A class containing lots of helper methods and state to build frames. This 26 allows test cases to easily build correct HTTP/2 frames to feed to 27 hyper-h2. 28 """ 29 def __init__(self): 30 self.encoder = Encoder() 31 32 def refresh_encoder(self): 33 self.encoder = Encoder() 34 35 def preamble(self): 36 return b'PRI * HTTP/2.0\r\n\r\nSM\r\n\r\n' 37 38 def build_headers_frame(self, 39 headers, 40 flags=[], 41 stream_id=1, 42 **priority_kwargs): 43 """ 44 Builds a single valid headers frame out of the contained headers. 45 """ 46 f = HeadersFrame(stream_id) 47 f.data = self.encoder.encode(headers) 48 f.flags.add('END_HEADERS') 49 for flag in flags: 50 f.flags.add(flag) 51 52 for k, v in priority_kwargs.items(): 53 setattr(f, k, v) 54 55 return f 56 57 def build_continuation_frame(self, header_block, flags=[], stream_id=1): 58 """ 59 Builds a single continuation frame out of the binary header block. 60 """ 61 f = ContinuationFrame(stream_id) 62 f.data = header_block 63 f.flags = set(flags) 64 65 return f 66 67 def build_data_frame(self, data, flags=None, stream_id=1, padding_len=0): 68 """ 69 Builds a single data frame out of a chunk of data. 70 """ 71 flags = set(flags) if flags is not None else set() 72 f = DataFrame(stream_id) 73 f.data = data 74 f.flags = flags 75 76 if padding_len: 77 flags.add('PADDED') 78 f.pad_length = padding_len 79 80 return f 81 82 def build_settings_frame(self, settings, ack=False): 83 """ 84 Builds a single settings frame. 85 """ 86 f = SettingsFrame(0) 87 if ack: 88 f.flags.add('ACK') 89 90 f.settings = settings 91 return f 92 93 def build_window_update_frame(self, stream_id, increment): 94 """ 95 Builds a single WindowUpdate frame. 96 """ 97 f = WindowUpdateFrame(stream_id) 98 f.window_increment = increment 99 return f 100 101 def build_ping_frame(self, ping_data, flags=None): 102 """ 103 Builds a single Ping frame. 104 """ 105 f = PingFrame(0) 106 f.opaque_data = ping_data 107 if flags: 108 f.flags = set(flags) 109 110 return f 111 112 def build_goaway_frame(self, 113 last_stream_id, 114 error_code=0, 115 additional_data=b''): 116 """ 117 Builds a single GOAWAY frame. 118 """ 119 f = GoAwayFrame(0) 120 f.error_code = error_code 121 f.last_stream_id = last_stream_id 122 f.additional_data = additional_data 123 return f 124 125 def build_rst_stream_frame(self, stream_id, error_code=0): 126 """ 127 Builds a single RST_STREAM frame. 128 """ 129 f = RstStreamFrame(stream_id) 130 f.error_code = error_code 131 return f 132 133 def build_push_promise_frame(self, 134 stream_id, 135 promised_stream_id, 136 headers, 137 flags=[]): 138 """ 139 Builds a single PUSH_PROMISE frame. 140 """ 141 f = PushPromiseFrame(stream_id) 142 f.promised_stream_id = promised_stream_id 143 f.data = self.encoder.encode(headers) 144 f.flags = set(flags) 145 f.flags.add('END_HEADERS') 146 return f 147 148 def build_priority_frame(self, 149 stream_id, 150 weight, 151 depends_on=0, 152 exclusive=False): 153 """ 154 Builds a single priority frame. 155 """ 156 f = PriorityFrame(stream_id) 157 f.depends_on = depends_on 158 f.stream_weight = weight 159 f.exclusive = exclusive 160 return f 161 162 def build_alt_svc_frame(self, stream_id, origin, field): 163 """ 164 Builds a single ALTSVC frame. 165 """ 166 f = AltSvcFrame(stream_id) 167 f.origin = origin 168 f.field = field 169 return f 170 171 def change_table_size(self, new_size): 172 """ 173 Causes the encoder to send a dynamic size update in the next header 174 block it sends. 175 """ 176 self.encoder.header_table_size = new_size