def test_proxy_headers(self, ClientRequestMock): req = ClientRequest( 'GET', URL('http://www.python.org'), proxy=URL('http://proxy.example.com'), proxy_headers={'Foo': 'Bar'}, loop=self.loop) self.assertEqual(str(req.proxy), 'http://proxy.example.com') # mock all the things! connector = aiohttp.TCPConnector(loop=self.loop) connector._resolve_host = make_mocked_coro([mock.MagicMock()]) proto = mock.Mock(**{ 'transport.get_extra_info.return_value': False, }) self.loop.create_connection = make_mocked_coro( (proto.transport, proto)) conn = self.loop.run_until_complete(connector.connect(req)) self.assertEqual(req.url, URL('http://www.python.org')) self.assertIs(conn._protocol, proto) self.assertIs(conn.transport, proto.transport) ClientRequestMock.assert_called_with( 'GET', URL('http://proxy.example.com'), auth=None, headers={'Host': 'www.python.org', 'Foo': 'Bar'}, loop=self.loop, ssl=None)
def test_connect_req_verify_ssl_false(self, ClientRequestMock): req = ClientRequest( 'GET', URL('https://www.python.org'), proxy=URL('http://proxy.example.com'), loop=self.loop, ssl=False, ) proto = mock.Mock() connector = aiohttp.TCPConnector(loop=self.loop) connector._create_proxy_connection = mock.MagicMock( side_effect=connector._create_proxy_connection) connector._create_direct_connection = mock.MagicMock( side_effect=connector._create_direct_connection) connector._resolve_host = make_mocked_coro([mock.MagicMock()]) self.loop.create_connection = make_mocked_coro( (proto.transport, proto)) self.loop.run_until_complete(connector.connect(req)) connector._create_proxy_connection.assert_called_with( req, traces=None) ((proxy_req,), _) = connector._create_direct_connection.call_args proxy_req.send.assert_called_with(mock.ANY)
def test__run_ok_max_requests_exceeded(worker, loop): worker.ppid = 1 worker.alive = True worker.servers = {} sock = mock.Mock() sock.cfg_addr = ('localhost', 8080) worker.sockets = [sock] worker.wsgi = mock.Mock() worker.close = make_mocked_coro(None) worker.log = mock.Mock() worker.loop = loop loop.create_server = make_mocked_coro(sock) worker.wsgi.make_handler.return_value.requests_count = 15 worker.cfg.access_log_format = ACCEPTABLE_LOG_FORMAT worker.cfg.max_requests = 10 worker.cfg.is_ssl = True ssl_context = mock.Mock() with mock.patch('ssl.SSLContext', return_value=ssl_context): with mock.patch('aiohttp.worker.asyncio') as m_asyncio: m_asyncio.sleep = mock.Mock( wraps=asyncio.coroutine(lambda *a, **kw: None)) yield from worker._run() worker.notify.assert_called_with() worker.log.info.assert_called_with("Max requests, shutting down: %s", worker) args, kwargs = loop.create_server.call_args assert 'ssl' in kwargs ctx = kwargs['ssl'] assert ctx is ssl_context
def test_connect(self, ClientRequestMock): req = ClientRequest('GET', URL('http://www.python.org'), loop=self.loop) self.assertEqual(req.url.path, '/') connector = aiohttp.ProxyConnector(URL('http://proxy.example.com'), loop=self.loop) self.assertIs(self.loop, connector._loop) connector._resolve_host = make_mocked_coro([mock.MagicMock()]) tr, proto = mock.Mock(), mock.Mock() self.loop.create_connection = make_mocked_coro((tr, proto)) conn = self.loop.run_until_complete(connector.connect(req)) self.assertEqual(req.url, URL('http://www.python.org')) self.assertIs(conn._transport, tr) self.assertIs(conn._protocol, proto) # resolve_host.assert_called_once_with('proxy.example.com', 80) tr.get_extra_info.assert_called_once_with('sslcontext') ClientRequestMock.assert_called_with( 'GET', URL('http://proxy.example.com'), auth=None, headers={'Host': 'www.python.org'}, loop=self.loop) conn.close()
def test_https_connect_resp_start_error(self, ClientRequestMock): proxy_req = ClientRequest('GET', URL('http://proxy.example.com'), loop=self.loop) ClientRequestMock.return_value = proxy_req proxy_resp = ClientResponse('get', URL('http://proxy.example.com')) proxy_resp._loop = self.loop proxy_req.send = send_mock = mock.Mock() send_mock.return_value = proxy_resp proxy_resp.start = make_mocked_coro( raise_exception=OSError("error message")) connector = aiohttp.TCPConnector(loop=self.loop) connector._resolve_host = make_mocked_coro( [{'hostname': 'hostname', 'host': '127.0.0.1', 'port': 80, 'family': socket.AF_INET, 'proto': 0, 'flags': 0}]) tr, proto = mock.Mock(), mock.Mock() tr.get_extra_info.return_value = None self.loop.create_connection = make_mocked_coro((tr, proto)) req = ClientRequest( 'GET', URL('https://www.python.org'), proxy=URL('http://proxy.example.com'), loop=self.loop, ) with self.assertRaisesRegex(OSError, "error message"): self.loop.run_until_complete(connector._create_connection(req))
def test_connect_env_var_not_overwriting(self, ClientRequestMock): req = ClientRequest( 'GET', URL('http://www.python.org'), proxy_from_env=True, proxy=URL('http://proxy.example.com'), loop=self.loop ) self.assertEqual(str(req.proxy), 'http://proxy.example.com') # mock all the things! connector = aiohttp.TCPConnector(loop=self.loop) connector._resolve_host = make_mocked_coro([mock.MagicMock()]) proto = mock.Mock() self.loop.create_connection = make_mocked_coro( (proto.transport, proto)) conn = self.loop.run_until_complete(connector.connect(req)) self.assertEqual(req.url, URL('http://www.python.org')) self.assertIs(conn._protocol, proto) self.assertIs(conn.transport, proto.transport) ClientRequestMock.assert_called_with( 'GET', URL('http://proxy.example.com'), auth=None, headers={'Host': 'www.python.org'}, loop=self.loop) self.assertIn('http', getproxies())
def test_auth_from_url(self, ClientRequestMock): proxy_req = ClientRequest('GET', URL('http://*****:*****@proxy.example.com'), loop=self.loop) ClientRequestMock.return_value = proxy_req self.assertIn('AUTHORIZATION', proxy_req.headers) self.assertNotIn('PROXY-AUTHORIZATION', proxy_req.headers) connector = aiohttp.TCPConnector(loop=self.loop) connector._resolve_host = make_mocked_coro([mock.MagicMock()]) tr, proto = mock.Mock(), mock.Mock() self.loop.create_connection = make_mocked_coro((tr, proto)) req = ClientRequest( 'GET', URL('http://www.python.org'), proxy=URL('http://*****:*****@proxy.example.com'), loop=self.loop, ) self.assertNotIn('AUTHORIZATION', req.headers) self.assertNotIn('PROXY-AUTHORIZATION', req.headers) conn = self.loop.run_until_complete(connector.connect(req)) self.assertEqual(req.url, URL('http://www.python.org')) self.assertNotIn('AUTHORIZATION', req.headers) self.assertIn('PROXY-AUTHORIZATION', req.headers) self.assertNotIn('AUTHORIZATION', proxy_req.headers) self.assertNotIn('PROXY-AUTHORIZATION', proxy_req.headers) ClientRequestMock.assert_called_with( 'GET', URL('http://*****:*****@proxy.example.com'), auth=None, loop=mock.ANY, headers=mock.ANY) conn.close()
def test_https_connect(self, ClientRequestMock): proxy_req = ClientRequest('GET', URL('http://proxy.example.com'), loop=self.loop) ClientRequestMock.return_value = proxy_req proxy_resp = ClientResponse('get', URL('http://proxy.example.com')) proxy_resp._loop = self.loop proxy_req.send = send_mock = mock.Mock() send_mock.return_value = proxy_resp proxy_resp.start = make_mocked_coro(mock.Mock(status=200)) connector = aiohttp.TCPConnector(loop=self.loop) connector._resolve_host = make_mocked_coro( [{'hostname': 'hostname', 'host': '127.0.0.1', 'port': 80, 'family': socket.AF_INET, 'proto': 0, 'flags': 0}]) tr, proto = mock.Mock(), mock.Mock() self.loop.create_connection = make_mocked_coro((tr, proto)) req = ClientRequest( 'GET', URL('https://www.python.org'), proxy=URL('http://proxy.example.com'), loop=self.loop, ) self.loop.run_until_complete(connector._create_connection(req)) self.assertEqual(req.url.path, '/') self.assertEqual(proxy_req.method, 'CONNECT') self.assertEqual(proxy_req.url, URL('https://www.python.org')) tr.pause_reading.assert_called_once_with() tr.get_extra_info.assert_called_with('socket', default=None) self.loop.run_until_complete(proxy_req.close()) proxy_resp.close() self.loop.run_until_complete(req.close())
def maker(method, path, query_params={}, headers=None, match_info=None, loop=None): path = URL(path) if query_params: path = path.with_query(query_params) if headers is None: headers = CIMultiDict( { "HOST": "server.example.com", "UPGRADE": "websocket", "CONNECTION": "Upgrade", "SEC-WEBSOCKET-KEY": "dGhlIHNhbXBsZSBub25jZQ==", "ORIGIN": "http://example.com", "SEC-WEBSOCKET-PROTOCOL": "chat, superchat", "SEC-WEBSOCKET-VERSION": "13", } ) writer = mock.Mock() writer.write_headers = make_mocked_coro(None) writer.write = make_mocked_coro(None) writer.drain = make_mocked_coro(None) transport = mock.Mock() transport._drain_helper = make_mocked_coro() ret = make_mocked_request(method, str(path), headers, writer=writer, loop=loop) if match_info is None: match_info = UrlMappingMatchInfo({}, mock.Mock()) match_info.add_app(app) ret._match_info = match_info return ret
async def make_app(): nonlocal startup_handler, cleanup_handler app = web.Application() startup_handler = make_mocked_coro() app.on_startup.append(startup_handler) cleanup_handler = make_mocked_coro() app.on_cleanup.append(cleanup_handler) return app
async def test_release_when_stream_at_eof(self): wrapper = MultipartResponseWrapper(mock.Mock(), mock.Mock()) wrapper.resp.release = make_mocked_coro(None) wrapper.stream.next = make_mocked_coro(b'') wrapper.stream.at_eof.return_value = True await wrapper.next() assert wrapper.stream.next.called assert wrapper.resp.release.called
async def test_app_register_on_finish() -> None: app = web.Application() cb1 = make_mocked_coro(None) cb2 = make_mocked_coro(None) app.on_cleanup.append(cb1) app.on_cleanup.append(cb2) app.freeze() await app.cleanup() cb1.assert_called_once_with(app) cb2.assert_called_once_with(app)
def patched_loop(loop): skip_if_no_dict(loop) server = mock.Mock() server.wait_closed = make_mocked_coro(None) loop.create_server = make_mocked_coro(server) unix_server = mock.Mock() unix_server.wait_closed = make_mocked_coro(None) loop.create_unix_server = make_mocked_coro(unix_server) asyncio.set_event_loop(loop) return loop
def test_https_auth(self, ClientRequestMock) -> None: proxy_req = ClientRequest('GET', URL('http://proxy.example.com'), auth=aiohttp.helpers.BasicAuth('user', 'pass'), loop=self.loop) ClientRequestMock.return_value = proxy_req proxy_resp = ClientResponse('get', URL('http://proxy.example.com'), request_info=mock.Mock(), writer=mock.Mock(), continue100=None, timer=TimerNoop(), traces=[], loop=self.loop, session=mock.Mock()) proxy_req.send = make_mocked_coro(proxy_resp) proxy_resp.start = make_mocked_coro(mock.Mock(status=200)) async def make_conn(): return aiohttp.TCPConnector() connector = self.loop.run_until_complete(make_conn()) connector._resolve_host = make_mocked_coro( [{'hostname': 'hostname', 'host': '127.0.0.1', 'port': 80, 'family': socket.AF_INET, 'proto': 0, 'flags': 0}]) tr, proto = mock.Mock(), mock.Mock() self.loop.create_connection = make_mocked_coro((tr, proto)) self.assertIn('AUTHORIZATION', proxy_req.headers) self.assertNotIn('PROXY-AUTHORIZATION', proxy_req.headers) req = ClientRequest( 'GET', URL('https://www.python.org'), proxy=URL('http://proxy.example.com'), loop=self.loop ) self.assertNotIn('AUTHORIZATION', req.headers) self.assertNotIn('PROXY-AUTHORIZATION', req.headers) self.loop.run_until_complete( connector._create_connection(req, None, aiohttp.ClientTimeout())) self.assertEqual(req.url.path, '/') self.assertNotIn('AUTHORIZATION', req.headers) self.assertNotIn('PROXY-AUTHORIZATION', req.headers) self.assertNotIn('AUTHORIZATION', proxy_req.headers) self.assertIn('PROXY-AUTHORIZATION', proxy_req.headers) connector._resolve_host.assert_called_with( 'proxy.example.com', 80, traces=mock.ANY) self.loop.run_until_complete(proxy_req.close()) proxy_resp.close() self.loop.run_until_complete(req.close())
async def test_consecutive_write_eof(): writer = mock.Mock() writer.write_eof = make_mocked_coro() writer.write_headers = make_mocked_coro() req = make_request('GET', '/', writer=writer) data = b'data' resp = Response(body=data) await resp.prepare(req) await resp.write_eof() await resp.write_eof() writer.write_eof.assert_called_once_with(data)
async def test_resolve_host_fail(loop, remote_resolve): tr, proto = mock.Mock(name='transport'), mock.Mock(name='protocol') with mock.patch('aiosocks.connector.create_connection', make_mocked_coro((tr, proto))): req = ProxyClientRequest( 'GET', URL('http://python.org'), loop=loop, proxy=URL('socks5://proxy.example')) connector = ProxyConnector(loop=loop, remote_resolve=remote_resolve) connector._resolve_host = make_mocked_coro(raise_exception=OSError()) with pytest.raises(aiohttp.ClientConnectorError): await connector.connect(req)
async def test_proxy_connect_http(loop): tr, proto = mock.Mock(name='transport'), mock.Mock(name='protocol') loop_mock = mock.Mock() loop_mock.getaddrinfo = make_mocked_coro([ [0, 0, 0, 0, ['127.0.0.1', 1080]]]) loop_mock.create_connection = make_mocked_coro((tr, proto)) req = ProxyClientRequest( 'GET', URL('http://python.org'), loop=loop, proxy=URL('http://127.0.0.1')) connector = ProxyConnector(loop=loop_mock) await connector.connect(req)
def test_startup_cleanup_signals_even_on_failure(patched_loop): patched_loop.create_server = mock.Mock(side_effect=RuntimeError()) app = web.Application() startup_handler = make_mocked_coro() app.on_startup.append(startup_handler) cleanup_handler = make_mocked_coro() app.on_cleanup.append(cleanup_handler) with pytest.raises(RuntimeError): web.run_app(app, print=stopper(patched_loop)) startup_handler.assert_called_once_with(app) cleanup_handler.assert_called_once_with(app)
async def test_proxy_negotiate_fail(loop): loop_mock = mock.Mock() loop_mock.getaddrinfo = make_mocked_coro( [[0, 0, 0, 0, ['127.0.0.1', 1080]]]) with mock.patch('aiosocks.connector.create_connection', make_mocked_coro(raise_exception=aiosocks.SocksError())): req = ProxyClientRequest( 'GET', URL('http://python.org'), loop=loop, proxy=URL('socks5://127.0.0.1')) connector = ProxyConnector(loop=loop_mock) with pytest.raises(aiosocks.SocksError): await connector.connect(req)
def test_https_connect_pass_ssl_context(self, ClientRequestMock) -> None: proxy_req = ClientRequest('GET', URL('http://proxy.example.com'), loop=self.loop) ClientRequestMock.return_value = proxy_req proxy_resp = ClientResponse('get', URL('http://proxy.example.com'), request_info=mock.Mock(), writer=mock.Mock(), continue100=None, timer=TimerNoop(), traces=[], loop=self.loop, session=mock.Mock()) proxy_req.send = make_mocked_coro(proxy_resp) proxy_resp.start = make_mocked_coro(mock.Mock(status=200)) async def make_conn(): return aiohttp.TCPConnector() connector = self.loop.run_until_complete(make_conn()) connector._resolve_host = make_mocked_coro( [{'hostname': 'hostname', 'host': '127.0.0.1', 'port': 80, 'family': socket.AF_INET, 'proto': 0, 'flags': 0}]) tr, proto = mock.Mock(), mock.Mock() self.loop.create_connection = make_mocked_coro((tr, proto)) req = ClientRequest( 'GET', URL('https://www.python.org'), proxy=URL('http://proxy.example.com'), loop=self.loop, ) self.loop.run_until_complete(connector._create_connection( req, None, aiohttp.ClientTimeout())) self.loop.create_connection.assert_called_with( mock.ANY, ssl=connector._make_ssl_context(True), sock=mock.ANY, server_hostname='www.python.org') self.assertEqual(req.url.path, '/') self.assertEqual(proxy_req.method, 'CONNECT') self.assertEqual(proxy_req.url, URL('https://www.python.org')) tr.close.assert_called_once_with() tr.get_extra_info.assert_called_with('socket', default=None) self.loop.run_until_complete(proxy_req.close()) proxy_resp.close() self.loop.run_until_complete(req.close())
def test_https_connect_ssl_error(self, ClientRequestMock) -> None: proxy_req = ClientRequest('GET', URL('http://proxy.example.com'), loop=self.loop) ClientRequestMock.return_value = proxy_req proxy_resp = ClientResponse('get', URL('http://proxy.example.com'), request_info=mock.Mock(), writer=mock.Mock(), continue100=None, timer=TimerNoop(), traces=[], loop=self.loop, session=mock.Mock()) proxy_req.send = make_mocked_coro(proxy_resp) proxy_resp.start = make_mocked_coro(mock.Mock(status=200)) async def make_conn(): return aiohttp.TCPConnector() connector = self.loop.run_until_complete(make_conn()) connector._resolve_host = make_mocked_coro( [{'hostname': 'hostname', 'host': '127.0.0.1', 'port': 80, 'family': socket.AF_INET, 'proto': 0, 'flags': 0}]) seq = 0 @asyncio.coroutine def create_connection(*args, **kwargs): nonlocal seq seq += 1 # connection to http://proxy.example.com if seq == 1: return mock.Mock(), mock.Mock() # connection to https://www.python.org elif seq == 2: raise ssl.SSLError else: assert False self.loop.create_connection = create_connection req = ClientRequest( 'GET', URL('https://www.python.org'), proxy=URL('http://proxy.example.com'), loop=self.loop, ) with self.assertRaises(aiohttp.ClientConnectorSSLError): self.loop.run_until_complete(connector._create_connection( req, None, aiohttp.ClientTimeout()))
async def test_negotiate_fail(): addr = aiosocks.Socks5Addr('localhost') auth = aiosocks.Socks5Auth('usr', 'pwd') dst = ('python.org', 80) loop_mock = mock.Mock() loop_mock.create_connection = make_mocked_coro((mock.Mock(), mock.Mock())) with mock.patch('aiosocks.asyncio.Future') as future_mock: future_mock.side_effect = make_mocked_coro( raise_exception=aiosocks.SocksError()) with pytest.raises(aiosocks.SocksError): await aiosocks.create_connection( None, addr, auth, dst, loop=loop_mock)
async def test_cannot_write_eof_twice(): resp = StreamResponse() writer = mock.Mock() resp_impl = await resp.prepare(make_request('GET', '/')) resp_impl.write = make_mocked_coro(None) resp_impl.write_eof = make_mocked_coro(None) await resp.write(b'data') assert resp_impl.write.called await resp.write_eof() resp_impl.write.reset_mock() await resp.write_eof() assert not writer.write.called
async def test_process_release_acquire_and_remote_closed(make_transport): transp = make_transport() transp.session.interrupted = False transp.manager.acquire = make_mocked_coro() transp.manager.release = make_mocked_coro() resp = await transp.process() await transp.manager.clear() assert resp.status == 101 assert resp.headers.get("upgrade", "").lower() == "websocket" assert resp.headers.get("connection", "").lower() == "upgrade" transp.session._remote_closed.assert_called_once_with() assert transp.manager.acquire.called assert transp.manager.release.called
def test_run_app_http(patched_loop): app = web.Application() startup_handler = make_mocked_coro() app.on_startup.append(startup_handler) cleanup_handler = make_mocked_coro() app.on_cleanup.append(cleanup_handler) web.run_app(app, print=stopper(patched_loop)) patched_loop.create_server.assert_called_with(mock.ANY, '0.0.0.0', 8080, ssl=None, backlog=128, reuse_address=None, reuse_port=None) startup_handler.assert_called_once_with(app) cleanup_handler.assert_called_once_with(app)
def test_using_gzip_if_header_present_and_file_available(loop): request = make_mocked_request( 'GET', URL('http://python.org/logo.png'), headers={ hdrs.ACCEPT_ENCODING: 'gzip' } ) gz_filepath = mock.Mock() gz_filepath.open = mock.mock_open() gz_filepath.is_file.return_value = True gz_filepath.stat.return_value = mock.MagicMock() gz_filepath.stat.st_size = 1024 filepath = mock.Mock() filepath.name = 'logo.png' filepath.open = mock.mock_open() filepath.with_name.return_value = gz_filepath file_sender = FileSender() file_sender._sendfile = make_mocked_coro(None) loop.run_until_complete(file_sender.send(request, filepath)) assert not filepath.open.called assert gz_filepath.open.called
async def test_ws_connect_deflate_per_message(loop, ws_key, key_data) -> None: resp = mock.Mock() resp.status = 101 resp.headers = { hdrs.UPGRADE: hdrs.WEBSOCKET, hdrs.CONNECTION: hdrs.UPGRADE, hdrs.SEC_WEBSOCKET_ACCEPT: ws_key, hdrs.SEC_WEBSOCKET_EXTENSIONS: 'permessage-deflate', } with mock.patch('aiohttp.client.WebSocketWriter') as WebSocketWriter: with mock.patch('aiohttp.client.os') as m_os: with mock.patch('aiohttp.client.ClientSession.request') as m_req: m_os.urandom.return_value = key_data m_req.return_value = loop.create_future() m_req.return_value.set_result(resp) writer = WebSocketWriter.return_value = mock.Mock() send = writer.send = make_mocked_coro() session = aiohttp.ClientSession(loop=loop) resp = await session.ws_connect('http://test.org') await resp.send_str('string', compress=-1) send.assert_called_with('string', binary=False, compress=-1) await resp.send_bytes(b'bytes', compress=15) send.assert_called_with(b'bytes', binary=True, compress=15) await resp.send_json([{}], compress=-9) send.assert_called_with('[{}]', binary=False, compress=-9) await session.close()
async def test_reader_read_exception(ws_key, key_data, loop) -> None: hresp = mock.Mock() hresp.status = 101 hresp.headers = { hdrs.UPGRADE: hdrs.WEBSOCKET, hdrs.CONNECTION: hdrs.UPGRADE, hdrs.SEC_WEBSOCKET_ACCEPT: ws_key, } with mock.patch('aiohttp.client.WebSocketWriter') as WebSocketWriter: with mock.patch('aiohttp.client.os') as m_os: with mock.patch('aiohttp.client.ClientSession.request') as m_req: m_os.urandom.return_value = key_data m_req.return_value = loop.create_future() m_req.return_value.set_result(hresp) writer = mock.Mock() WebSocketWriter.return_value = writer writer.close = make_mocked_coro() session = aiohttp.ClientSession(loop=loop) resp = await session.ws_connect('http://test.org') exc = ValueError() resp._reader.set_exception(exc) msg = await resp.receive() assert msg.type == aiohttp.WSMsgType.ERROR assert resp.exception() is exc await session.close()
async def test_next(self): wrapper = MultipartResponseWrapper(mock.Mock(), mock.Mock()) wrapper.stream.next = make_mocked_coro(b'') wrapper.stream.at_eof.return_value = False await wrapper.next() assert wrapper.stream.next.called
async def test_close_exc(loop, ws_key, key_data): resp = mock.Mock() resp.status = 101 resp.headers = { hdrs.UPGRADE: hdrs.WEBSOCKET, hdrs.CONNECTION: hdrs.UPGRADE, hdrs.SEC_WEBSOCKET_ACCEPT: ws_key, } with mock.patch('aiohttp.client.WebSocketWriter') as WebSocketWriter: with mock.patch('aiohttp.client.os') as m_os: with mock.patch('aiohttp.client.ClientSession.get') as m_req: m_os.urandom.return_value = key_data m_req.return_value = loop.create_future() m_req.return_value.set_result(resp) writer = mock.Mock() WebSocketWriter.return_value = writer writer.close = make_mocked_coro() session = aiohttp.ClientSession(loop=loop) resp = await session.ws_connect('http://test.org') assert not resp.closed exc = ValueError() resp._reader.set_exception(exc) await resp.close() assert resp.closed assert resp.exception() is exc await session.close()
async def test_chunked_explicit(loop, conn): req = ClientRequest('post', URL('http://python.org/'), chunked=True, loop=loop) with mock.patch('aiohttp.client_reqrep.StreamWriter') as m_writer: m_writer.return_value.write_headers = make_mocked_coro() resp = await req.send(conn) assert 'chunked' == req.headers['TRANSFER-ENCODING'] m_writer.return_value.enable_chunking.assert_called_with() await req.close() resp.close()
async def test_shutdown_timeout(loop): manager = web.Server(serve, loop=loop) handler = mock.Mock() handler.shutdown = make_mocked_coro(mock.Mock()) transport = mock.Mock() manager.connection_made(handler, transport) await manager.shutdown(timeout=0.1) manager.connection_lost(handler, None) assert manager.connections == [] handler.shutdown.assert_called_with(0.1)
def function2045(self): var1894 = aiohttp.TCPConnector(loop=self.attribute501) var1894._resolve_host = make_mocked_coro( raise_exception=OSError('dont take it serious')) var4452 = ClientRequest('GET', URL('http://www.python.org'), proxy=URL('http://proxy.example.com'), loop=self.attribute501) var1607 = dict(var4452.headers) with self.assertRaises(aiohttp.ClientConnectorError): self.attribute501.run_until_complete(var1894.connect(var4452)) self.assertEqual(var4452.url.path, '/') self.assertEqual(dict(var4452.headers), var1607)
async def test_streaming_send(make_transport): trans = make_transport() resp = trans.response = mock.Mock() resp.write = make_mocked_coro(None) stop = await trans.send('text data') assert not stop assert trans.size == len(b'text data\n') resp.write.assert_called_with(b'text data\n') trans.maxsize = 1 stop = await trans.send('text data') assert stop
async def test_reload(method, mocker): if method: mocker.patch("sys.executable", "python" if method == "execl" else "run") else: mocker.patch("sys.executable", "foobar") execl = mocker.patch("os.execl") execv = mocker.patch("os.execv") app = mocker.Mock() setattr(app, "on_shutdown", [make_mocked_coro(), make_mocked_coro()]) if method: await virtool.utils.reload(app) else: with pytest.raises(SystemError): await virtool.utils.reload(app) assert execl.called is (method == "execl") assert execv.called is (method == "execv")
async def test_after_headers_doesnt_request_random_filename_if_redirect_is_already_in_knowledge_base( self): self.fake_engine.perform_high_priority = make_mocked_coro() kb = KnowledgeBase() self.heuristic.set_kb(kb) kb.redirects["%s/wp-admin/" % self.host] = "%s/somewhere" % self.host entry = self.create_redirected_request("/wp-admin/admin.php", redirected_to="/wp-login.php") await self.heuristic.after_headers(entry) self.fake_engine.perform_high_priority.assert_not_called()
async def test_after_headers_transform_relative_location_to_absolute_location( self): initial_request = self.create_redirected_request( "/admin/resource.php", redirected_to="/admin/login.php", relative=True) returned_entry = self.create_redirected_request( "/admin/uuid", redirected_to="/admin/login.php") self.fake_engine.perform_high_priority = make_mocked_coro( return_value=returned_entry) with self.assertRaises(RejectRequest): await self.heuristic.after_headers(initial_request)
def test_https_connect_runtime_error(self, ClientRequestMock): proxy_req = ClientRequest('GET', URL('http://proxy.example.com'), loop=self.loop) ClientRequestMock.return_value = proxy_req proxy_resp = ClientResponse('get', URL('http://proxy.example.com')) proxy_resp._loop = self.loop proxy_req.send = send_mock = mock.Mock() send_mock.return_value = proxy_resp proxy_resp.start = make_mocked_coro(mock.Mock(status=200)) connector = aiohttp.TCPConnector(loop=self.loop) connector._resolve_host = make_mocked_coro([{ 'hostname': 'hostname', 'host': '127.0.0.1', 'port': 80, 'family': socket.AF_INET, 'proto': 0, 'flags': 0 }]) tr, proto = mock.Mock(), mock.Mock() tr.get_extra_info.return_value = None self.loop.create_connection = make_mocked_coro((tr, proto)) req = ClientRequest( 'GET', URL('https://www.python.org'), proxy=URL('http://proxy.example.com'), loop=self.loop, ) with self.assertRaisesRegex( RuntimeError, "Transport does not expose socket instance"): self.loop.run_until_complete(connector._create_connection(req)) self.loop.run_until_complete(proxy_req.close()) proxy_resp.close() self.loop.run_until_complete(req.close())
def maker(method, path, query_params={}, headers=None, match_info=None, loop=None): path = URL(path) if query_params: path = path.with_query(query_params) if headers is None: headers = CIMultiDict({ 'HOST': 'server.example.com', 'UPGRADE': 'websocket', 'CONNECTION': 'Upgrade', 'SEC-WEBSOCKET-KEY': 'dGhlIHNhbXBsZSBub25jZQ==', 'ORIGIN': 'http://example.com', 'SEC-WEBSOCKET-PROTOCOL': 'chat, superchat', 'SEC-WEBSOCKET-VERSION': '13' }) writer = mock.Mock() writer.write_headers = make_mocked_coro(None) writer.write = make_mocked_coro(None) writer.drain = make_mocked_coro(None) transport = mock.Mock() transport._drain_helper = make_mocked_coro() ret = make_mocked_request(method, str(path), headers, writer=writer, loop=loop) if match_info is None: match_info = UrlMappingMatchInfo({}, mock.Mock()) match_info.add_app(app) ret._match_info = match_info return ret
def test_auth_from_url(self, ClientRequestMock): proxy_req = ClientRequest('GET', URL('http://*****:*****@proxy.example.com'), loop=self.loop) ClientRequestMock.return_value = proxy_req self.assertIn('AUTHORIZATION', proxy_req.headers) self.assertNotIn('PROXY-AUTHORIZATION', proxy_req.headers) connector = aiohttp.TCPConnector(loop=self.loop) connector._resolve_host = make_mocked_coro([mock.MagicMock()]) tr, proto = mock.Mock(), mock.Mock() self.loop.create_connection = make_mocked_coro((tr, proto)) req = ClientRequest( 'GET', URL('http://www.python.org'), proxy=URL('http://*****:*****@proxy.example.com'), loop=self.loop, ) self.assertNotIn('AUTHORIZATION', req.headers) self.assertNotIn('PROXY-AUTHORIZATION', req.headers) conn = self.loop.run_until_complete(connector.connect(req)) self.assertEqual(req.url, URL('http://www.python.org')) self.assertNotIn('AUTHORIZATION', req.headers) self.assertIn('PROXY-AUTHORIZATION', req.headers) self.assertNotIn('AUTHORIZATION', proxy_req.headers) self.assertNotIn('PROXY-AUTHORIZATION', proxy_req.headers) ClientRequestMock.assert_called_with( 'GET', URL('http://*****:*****@proxy.example.com'), auth=None, loop=mock.ANY, headers=mock.ANY, ssl=None) conn.close()
async def test_shutdown_no_timeout(loop): app = web.Application() manager = app.make_handler(loop=loop) handler = mock.Mock() handler.shutdown = make_mocked_coro(mock.Mock()) transport = mock.Mock() manager.connection_made(handler, transport) await manager.shutdown() manager.connection_lost(handler, None) assert manager.connections == [] handler.shutdown.assert_called_with(None)
async def test_proxy_connect_http(loop): tr, proto = mock.Mock(name='transport'), mock.Mock(name='protocol') loop_mock = mock.Mock() loop_mock.getaddrinfo = make_mocked_coro( [[0, 0, 0, 0, ['127.0.0.1', 1080]]]) loop_mock.create_connection = make_mocked_coro((tr, proto)) loop_mock.create_task.return_value = asyncio.Task( make_mocked_coro([{ 'host': 'host', 'port': 80, 'family': 1, 'hostname': 'hostname', 'flags': 11, 'proto': 'proto' }])()) req = ProxyClientRequest('GET', URL('http://python.org'), loop=loop, proxy=URL('http://127.0.0.1')) connector = ProxyConnector(loop=loop_mock) await connector.connect(req, [], ClientTimeout())
async def test_sleep(self, mocker, dbi, test_blast_obj): """ Test that calling the `sleep()` method calls `asyncio.sleep` and increments the `interval` attribute by 5. """ m_sleep = mocker.patch("asyncio.sleep", make_mocked_coro()) await test_blast_obj.sleep() m_sleep.assert_called_with(3) assert test_blast_obj.interval == 8 await test_blast_obj.sleep() m_sleep.assert_called_with(8) assert test_blast_obj.interval == 13
async def test_send_with_per_message_deflate(make_request, mocker): req = make_request('GET', '/') ws = WebSocketResponse() await ws.prepare(req) writer_send = ws._writer.send = make_mocked_coro() await ws.send_str('string', compress=15) writer_send.assert_called_with('string', binary=False, compress=15) await ws.send_bytes(b'bytes', compress=0) writer_send.assert_called_with(b'bytes', binary=True, compress=0) await ws.send_json('[{}]', compress=9) writer_send.assert_called_with('"[{}]"', binary=False, compress=9)
async def test_prepare_calls_signal(): app = mock.Mock() sig = make_mocked_coro() on_response_prepare = signals.Signal(app) on_response_prepare.append(sig) req = make_request('GET', '/', app=app, on_response_prepare=on_response_prepare) resp = StreamResponse() await resp.prepare(req) sig.assert_called_with(req, resp)
async def test_prepare_calls_signal() -> None: app = mock.Mock() sig = make_mocked_coro() on_response_prepare = aiosignal.Signal(app) on_response_prepare.append(sig) req = make_request("GET", "/", app=app, on_response_prepare=on_response_prepare) resp = StreamResponse() await resp.prepare(req) sig.assert_called_with(req, resp)
def test_finish_connection_timeout(loop): app = web.Application() manager = app.make_handler(loop=loop) handler = mock.Mock() handler.shutdown = make_mocked_coro(mock.Mock()) transport = mock.Mock() manager.connection_made(handler, transport) yield from manager.finish_connections(timeout=0.1) manager.connection_lost(handler, None) assert manager.connections == [] handler.shutdown.assert_called_with(0.1)
async def test_base_make_ssl_proto(): loop_mock = mock.Mock() app_proto = mock.Mock() ssl_context = ssllib.create_default_context() proto = make_base(loop_mock, ap_factory=lambda: app_proto, ssl=ssl_context) proto.socks_request = make_mocked_coro((None, None)) proto._transport = mock.Mock() await proto.negotiate(None, None) mtr = loop_mock._make_ssl_transport assert mtr.called assert mtr.call_args[1]['sslcontext'] is ssl_context
async def test_paths_exists_fetch_generated_paths(self, loop): path_generator = MagicMock() path_generator.generate_paths.return_value = ["/", "/test", "/path"] fake_directory_fetcher = MagicMock() fake_directory_fetcher.fetch_paths = make_mocked_coro() tachyon.PathGenerator = MagicMock(return_value=path_generator) tachyon.DirectoryFetcher = MagicMock( return_value=fake_directory_fetcher) await tachyon.test_paths_exists(HammerTime(loop=loop), accumulator=self.accumulator) fake_directory_fetcher.fetch_paths.assert_called_once_with( path_generator.generate_paths.return_value)
async def test_perform_use_proxy_for_request(self, loop): asyncio.set_event_loop(loop) engine = AioHttpEngine(loop=loop, proxy="http://some.proxy.com/") engine.session.request = make_mocked_coro(return_value=FakeResponse()) entry = Entry.create("http://www.example.com/1") await engine.perform(entry, Heuristics()) engine.session.request.assert_called_once_with( method=entry.request.method, url="http://www.example.com/1", timeout=0.2, proxy="http://some.proxy.com/", allow_redirects=False)
async def test_connection_fail(): addr = aiosocksy.Socks5Addr('localhost') auth = aiosocksy.Socks5Auth('usr', 'pwd') dst = ('python.org', 80) loop_mock = mock.Mock() loop_mock.create_connection = make_mocked_coro(raise_exception=OSError()) with pytest.raises(aiosocksy.SocksConnectionError): await aiosocksy.create_connection(None, addr, auth, dst, loop=loop_mock)
def function2260(function742, arg841): var608 = mock.Mock() var608.wait_closed = make_mocked_coro(None) var3709 = mock.Mock() function742.servers = { srv: handler, } function742.log = mock.Mock() function742.loop = arg841 var3538 = function742.wsgi = mock.Mock() var3538.cleanup = make_mocked_coro(None) var3709.connections = [object()] var3709.shutdown.return_value = helpers.create_future(arg841) var3709.shutdown.return_value.set_result(1) var3538.shutdown.return_value = helpers.create_future(arg841) var3538.shutdown.return_value.set_result(None) yield from function742.close() var3538.shutdown.assert_called_with() var3538.cleanup.assert_called_with() var3709.shutdown.assert_called_with(timeout=95.0) var608.close.assert_called_with() assert (function742.servers is None) yield from function742.close()
async def test_content_encoding(loop: Any, conn: Any) -> None: req = ClientRequest("post", URL("http://python.org/"), data="foo", compress="deflate", loop=loop) with mock.patch("aiohttp.client_reqrep.StreamWriter") as m_writer: m_writer.return_value.write_headers = make_mocked_coro() resp = await req.send(conn) assert req.headers["TRANSFER-ENCODING"] == "chunked" assert req.headers["CONTENT-ENCODING"] == "deflate" m_writer.return_value.enable_compression.assert_called_with("deflate") await req.close() resp.close()
def test_receive_exc_in_reader(make_request, loop, reader): req = make_request('GET', '/') ws = WebSocketResponse() yield from ws.prepare(req) exc = ValueError() res = helpers.create_future(loop) res.set_exception(exc) reader.read = make_mocked_coro(res) msg = yield from ws.receive() assert msg.tp == WSMsgType.error assert msg.data is exc assert ws.exception() is exc
def function1340(self, arg735): var2739 = ClientRequest('GET', URL('http://www.python.org'), proxy=URL('http://proxy.example.com'), loop=self.attribute501) self.assertEqual(str(var2739.proxy), 'http://proxy.example.com') var3234 = aiohttp.TCPConnector(loop=self.attribute501) var3234._resolve_host = make_mocked_coro([mock.MagicMock()]) var4104 = mock.Mock() self.attribute501.create_connection = make_mocked_coro( (var4104.transport, var4104)) var1447 = self.attribute501.run_until_complete( var3234.connect(var2739)) self.assertEqual(var2739.url, URL('http://www.python.org')) self.assertIs(var1447._protocol, var4104) self.assertIs(var1447.transport, var4104.transport) arg735.assert_called_with('GET', URL('http://proxy.example.com'), auth=None, headers={ 'Host': 'www.python.org', }, loop=self.attribute501)
async def test_set_ssl_parameters_dont_verify_ssl(self, loop): engine = AioHttpEngine(loop=loop) engine.session.request = make_mocked_coro(return_value=FakeResponse()) entry = Entry.create("http://www.example.com/") engine.set_ssl_parameters(verify_ssl=False) await engine.perform(entry, Heuristics()) engine.session.request.assert_called_once_with( method=entry.request.method, url=entry.request.url, allow_redirects=False, timeout=0.2, ssl=False)
async def test_content_encoding_header(loop, conn) -> None: req = ClientRequest('post', URL('http://python.org/'), data='foo', headers={'Content-Encoding': 'deflate'}, loop=loop) with mock.patch('aiohttp.client_reqrep.StreamWriter') as m_writer: m_writer.return_value.write_headers = make_mocked_coro() resp = await req.send(conn) assert not m_writer.return_value.enable_compression.called assert not m_writer.return_value.enable_chunking.called await req.close() resp.close()
async def test_get_status(mocker, spawn_client): client = await spawn_client(authorize=True) mocker.patch("virtool.hmm.db.get_status", make_mocked_coro({ "id": "hmm", "updating": True })) resp = await client.get("/api/hmms/status") assert resp.status == 200 assert await resp.json() == {"id": "hmm", "updating": True}
async def test_send_with_per_message_deflate(make_request, mocker) -> None: req = make_request("GET", "/") ws = WebSocketResponse() await ws.prepare(req) writer_send = ws._writer.send = make_mocked_coro() await ws.send_str("string", compress=15) writer_send.assert_called_with("string", binary=False, compress=15) await ws.send_bytes(b"bytes", compress=0) writer_send.assert_called_with(b"bytes", binary=True, compress=0) await ws.send_json("[{}]", compress=9) writer_send.assert_called_with('"[{}]"', binary=False, compress=9)
def test_connect_env_var_https_ignored(self, ClientRequestMock): req = ClientRequest('GET', URL('http://www.python.org'), proxy_from_env=True, loop=self.loop) self.assertIsNone(req.proxy) # mock all the things! connector = aiohttp.TCPConnector(loop=self.loop) connector._resolve_host = make_mocked_coro([mock.MagicMock()]) proto = mock.Mock() self.loop.create_connection = make_mocked_coro( (proto.transport, proto)) conn = self.loop.run_until_complete(connector.connect(req)) self.assertEqual(req.url, URL('http://www.python.org')) self.assertIs(conn._protocol, proto) self.assertIs(conn.transport, proto.transport) # we patch only the connector, it should not be called ClientRequestMock.assert_not_called() self.assertIn('https', getproxies()) self.assertNotIn('http', getproxies())