|
14 | 14 | import unittest |
15 | 15 | from unittest import mock |
16 | 16 |
|
17 | | -from proxy.proxy import main, entry_point |
| 17 | +import requests |
| 18 | + |
| 19 | +from proxy.proxy import Proxy, main, entry_point |
18 | 20 | from proxy.common.utils import bytes_ |
19 | 21 | from proxy.core.work.fd import RemoteFdExecutor |
20 | 22 | from proxy.common.constants import ( # noqa: WPS450 |
@@ -368,3 +370,23 @@ def test_enable_ssh_tunnel( |
368 | 370 | mock_ssh_tunnel_listener.return_value.shutdown.assert_called_once() |
369 | 371 | # shutdown will internally call stop port forward |
370 | 372 | mock_ssh_tunnel_listener.return_value.stop_port_forward.assert_not_called() |
| 373 | + |
| 374 | + |
| 375 | +class TestProxyContextManager(unittest.TestCase): |
| 376 | + |
| 377 | + def test_proxy_context_manager(self) -> None: |
| 378 | + with Proxy(port=8888, num_acceptors=1): |
| 379 | + response = requests.get( |
| 380 | + 'http://httpbin.org/get', proxies={ |
| 381 | + 'http': 'http://127.0.0.1:8888', |
| 382 | + 'https': 'http://127.0.0.1:8888', |
| 383 | + }, |
| 384 | + ) |
| 385 | + self.assertEqual(response.status_code, 200) |
| 386 | + response = requests.get( |
| 387 | + 'https://httpbin.org/get', proxies={ |
| 388 | + 'http': 'http://127.0.0.1:8888', |
| 389 | + 'https': 'http://127.0.0.1:8888', |
| 390 | + }, |
| 391 | + ) |
| 392 | + self.assertEqual(response.status_code, 200) |
0 commit comments