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
|
"""Test using a proxy."""
import asyncio
import http.server
import socketserver
import threading
from urllib.request import urlopen
import pytest
import vcr
# Conditional imports
requests = pytest.importorskip("requests")
class Proxy(http.server.SimpleHTTPRequestHandler):
"""
Simple proxy server.
(Inspired by: http://effbot.org/librarybook/simplehttpserver.htm).
"""
def do_GET(self):
upstream_response = urlopen(self.path)
try:
status = upstream_response.status
headers = upstream_response.headers.items()
except AttributeError:
# In Python 2 the response is an addinfourl instance.
status = upstream_response.code
headers = upstream_response.info().items()
self.log_request(status)
self.send_response_only(status, upstream_response.msg)
for header in headers:
self.send_header(*header)
self.end_headers()
self.copyfile(upstream_response, self.wfile)
def do_CONNECT(self):
host, port = self.path.split(":")
asyncio.run(self._tunnel(host, port, self.connection))
async def _tunnel(self, host, port, client_sock):
target_r, target_w = await asyncio.open_connection(host=host, port=port)
self.send_response(http.HTTPStatus.OK)
self.end_headers()
source_r, source_w = await asyncio.open_connection(sock=client_sock)
async def channel(reader, writer):
while True:
data = await reader.read(1024)
if not data:
break
writer.write(data)
await writer.drain()
writer.close()
await writer.wait_closed()
await asyncio.gather(
channel(target_r, source_w),
channel(source_r, target_w),
)
@pytest.fixture(scope="session")
def proxy_server():
with socketserver.ThreadingTCPServer(("", 0), Proxy) as httpd:
proxy_process = threading.Thread(target=httpd.serve_forever)
proxy_process.start()
yield "http://{}:{}".format(*httpd.server_address)
httpd.shutdown()
proxy_process.join()
def test_use_proxy(tmpdir, httpbin, proxy_server):
"""Ensure that it works with a proxy."""
with vcr.use_cassette(str(tmpdir.join("proxy.yaml"))):
response = requests.get(httpbin.url, proxies={"http": proxy_server})
with vcr.use_cassette(str(tmpdir.join("proxy.yaml")), mode="none") as cassette:
cassette_response = requests.get(httpbin.url, proxies={"http": proxy_server})
assert cassette_response.headers == response.headers
assert cassette.play_count == 1
def test_use_https_proxy(tmpdir, httpbin_secure, proxy_server):
"""Ensure that it works with an HTTPS proxy."""
with vcr.use_cassette(str(tmpdir.join("proxy.yaml"))):
response = requests.get(httpbin_secure.url, proxies={"https": proxy_server})
with vcr.use_cassette(str(tmpdir.join("proxy.yaml")), mode="none") as cassette:
cassette_response = requests.get(
httpbin_secure.url,
proxies={"https": proxy_server},
)
assert cassette_response.headers == response.headers
assert cassette.play_count == 1
# The cassette URL points to httpbin, not the proxy
assert cassette.requests[0].url == httpbin_secure.url + "/"
|