File: test_proxy.py

package info (click to toggle)
vcr.py 8.1.1-1
  • links: PTS, VCS
  • area: main
  • in suites: forky, sid
  • size: 1,036 kB
  • sloc: python: 6,275; makefile: 188; sh: 1
file content (107 lines) | stat: -rw-r--r-- 3,396 bytes parent folder | download
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 + "/"