Rewrite test_proxymanager.py to be pytest-style

1 files changed, 29 insertions, 39 deletions

diff --git a/test/test_proxymanager.py b/test/test_proxymanager.py
index fa525c9f..ed218106 100644
--- a/test/test_proxymanager.py
+++ b/test/test_proxymanager.py
@@ -1,55 +1,45 @@
-import sys
+import pytest
 
 from urllib3.poolmanager import ProxyManager
 
-if sys.version_info >= (2, 7):
-    import unittest
-else:
-    import unittest2 as unittest
 
-
-class TestProxyManager(unittest.TestCase):
+class TestProxyManager(object):
     def test_proxy_headers(self):
-        p = ProxyManager('http://something:1234')
-        self.addCleanup(p.clear)
         url = 'http://pypi.python.org/test'
+        with ProxyManager('http://something:1234') as p:
 
-        # Verify default headers
-        default_headers = {'Accept': '*/*',
-                           'Host': 'pypi.python.org'}
-        headers = p._set_proxy_headers(url)
+            # Verify default headers
+            default_headers = {'Accept': '*/*',
+                               'Host': 'pypi.python.org'}
+            headers = p._set_proxy_headers(url)
 
-        self.assertEqual(headers, default_headers)
+            assert headers == default_headers
 
-        # Verify default headers don't overwrite provided headers
-        provided_headers = {'Accept': 'application/json',
-                            'custom': 'header',
-                            'Host': 'test.python.org'}
-        headers = p._set_proxy_headers(url, provided_headers)
+            # Verify default headers don't overwrite provided headers
+            provided_headers = {'Accept': 'application/json',
+                                'custom': 'header',
+                                'Host': 'test.python.org'}
+            headers = p._set_proxy_headers(url, provided_headers)
 
-        self.assertEqual(headers, provided_headers)
+            assert headers == provided_headers
 
-        # Verify proxy with nonstandard port
-        provided_headers = {'Accept': 'application/json'}
-        expected_headers = provided_headers.copy()
-        expected_headers.update({'Host': 'pypi.python.org:8080'})
-        url_with_port = 'http://pypi.python.org:8080/test'
-        headers = p._set_proxy_headers(url_with_port, provided_headers)
+            # Verify proxy with nonstandard port
+            provided_headers = {'Accept': 'application/json'}
+            expected_headers = provided_headers.copy()
+            expected_headers.update({'Host': 'pypi.python.org:8080'})
+            url_with_port = 'http://pypi.python.org:8080/test'
+            headers = p._set_proxy_headers(url_with_port, provided_headers)
 
-        self.assertEqual(headers, expected_headers)
+            assert headers == expected_headers
 
     def test_default_port(self):
-        p = ProxyManager('http://something')
-        self.addCleanup(p.clear)
-        self.assertEqual(p.proxy.port, 80)
-        p = ProxyManager('https://something')
-        self.addCleanup(p.clear)
-        self.assertEqual(p.proxy.port, 443)
+        with ProxyManager('http://something') as p:
+            assert p.proxy.port == 80
+        with ProxyManager('https://something') as p:
+            assert p.proxy.port == 443
 
     def test_invalid_scheme(self):
-        self.assertRaises(AssertionError, ProxyManager, 'invalid://host/p')
-        self.assertRaises(ValueError, ProxyManager, 'invalid://host/p')
-
-
-if __name__ == '__main__':
-    unittest.main()
+        with pytest.raises(AssertionError):
+            ProxyManager('invalid://host/p')
+        with pytest.raises(ValueError):
+            ProxyManager('invalid://host/p')