Skip to content

Commit

Permalink
fix bug with http not working with last commit
Browse files Browse the repository at this point in the history
  • Loading branch information
willforde committed Jan 24, 2018
1 parent 83deef9 commit 4bb2275
Show file tree
Hide file tree
Showing 2 changed files with 32 additions and 17 deletions.
31 changes: 20 additions & 11 deletions tests/test_urlquick.py
Original file line number Diff line number Diff line change
Expand Up @@ -680,17 +680,26 @@ def test_connect(self):
resp = cm.connect(req, 10, True)

self.assertIsInstance(resp, self.Response)
self.assertTrue("httpbin.org" in cm.request_handler["https"][1])
self.assertIsInstance(cm.request_handler["https"][1]["httpbin.org"], self.HTTPConnection)
self.assertTrue("httpbin.org" in cm.request_handler["https"])
self.assertIsInstance(cm.request_handler["https"]["httpbin.org"], self.HTTPConnection)

def test_connect_unverify(self):
req = urlquick.Request("GET", "https://httpbin.org/get", urlquick.CaseInsensitiveDict())
cm = urlquick.ConnectionManager()
resp = cm.connect(req, 10, False)

self.assertIsInstance(resp, self.Response)
self.assertTrue("httpbin.org" in cm.request_handler["https"][1])
self.assertIsInstance(cm.request_handler["https"][1]["httpbin.org"], self.HTTPConnection)
self.assertTrue("httpbin.org" in cm.request_handler["https"])
self.assertIsInstance(cm.request_handler["https"]["httpbin.org"], self.HTTPConnection)

def test_connect_unverify_http(self):
req = urlquick.Request("GET", "http://httpbin.org/get", urlquick.CaseInsensitiveDict())
cm = urlquick.ConnectionManager()
resp = cm.connect(req, 10, False)

self.assertIsInstance(resp, self.Response)
self.assertTrue("httpbin.org" in cm.request_handler["http"])
self.assertIsInstance(cm.request_handler["http"]["httpbin.org"], self.HTTPConnection)

def test_connect_will_close(self):
req = urlquick.Request("GET", "https://httpbin.org/get", urlquick.CaseInsensitiveDict())
Expand All @@ -700,33 +709,33 @@ def test_connect_will_close(self):
resp = cm.connect(req, 10, True)

self.assertIsInstance(resp, self.Response)
self.assertFalse("httpbin.org" in cm.request_handler["https"][1])
self.assertFalse("httpbin.org" in cm.request_handler["https"])
finally:
self.Response.will_close = False

def test_connect_reuse_good(self):
req = urlquick.Request("GET", "https://httpbin.org/get", urlquick.CaseInsensitiveDict())
cm = urlquick.ConnectionManager()
cm.request_handler["https"][1]["httpbin.org"] = self.HTTPConnection()
cm.request_handler["https"]["httpbin.org"] = self.HTTPConnection()
resp = cm.connect(req, 10, True)
self.assertIsInstance(resp, self.Response)

def test_connect_reuse_bad(self):
req = urlquick.Request("GET", "https://httpbin.org/get", urlquick.CaseInsensitiveDict())
cm = urlquick.ConnectionManager()
cm.request_handler["https"][1]["httpbin.org"] = self.HTTPConnection(fail=urlquick.HTTPException)
cm.request_handler["https"]["httpbin.org"] = self.HTTPConnection(fail=urlquick.HTTPException)
self.Response.will_close = True
try:
resp = cm.connect(req, 10, True)
self.assertIsInstance(resp, self.Response)
self.assertFalse("httpbin.org" in cm.request_handler["https"][1])
self.assertFalse("httpbin.org" in cm.request_handler["https"])
finally:
self.Response.will_close = False

def test_connect_reuse_ugly(self):
req = urlquick.Request("GET", "https://httpbin.org/get", urlquick.CaseInsensitiveDict())
cm = urlquick.ConnectionManager()
cm.request_handler["https"][1]["httpbin.org"] = self.HTTPConnection(fail=RuntimeError)
cm.request_handler["https"]["httpbin.org"] = self.HTTPConnection(fail=RuntimeError)
with self.assertRaises(RuntimeError):
cm.connect(req, 10, True)

Expand Down Expand Up @@ -772,9 +781,9 @@ def test_send_request_raises_httperror(self):

def test_close_connections(self):
cm = urlquick.ConnectionManager()
cm.request_handler["https"][1]["httpbin.org"] = self.HTTPConnection()
cm.request_handler["https"]["httpbin.org"] = self.HTTPConnection()
cm.close()
self.assertFalse(cm.request_handler["https"][1])
self.assertFalse(cm.request_handler["https"])


def create_resp(body=b"", headers=None, status=200, reason="OK"):
Expand Down
18 changes: 12 additions & 6 deletions urlquick.py
Original file line number Diff line number Diff line change
Expand Up @@ -504,7 +504,7 @@ def close(self):

class ConnectionManager(CacheAdapter):
def __init__(self):
self.request_handler = {"http": (HTTPConnection, {}), "https": (HTTPSConnection, {})}
self.request_handler = {"http": {}, "https": {}}
super(ConnectionManager, self).__init__()

def make_request(self, req, timeout, verify, max_age):
Expand All @@ -528,9 +528,10 @@ def make_request(self, req, timeout, verify, max_age):

def connect(self, req, timeout, verify):
# Fetch connection from pool and attempt to reuse if available
connection, pool = self.request_handler[req.type]
pool = self.request_handler[req.type]
if req.host in pool:
try:
# noinspection PyTypeChecker
return self.send_request(pool[req.host], req)
except Exception as e:
# Remove the connection from the pool as it's unusable
Expand All @@ -542,9 +543,14 @@ def connect(self, req, timeout, verify):
raise

# Create a new connection
# noinspection PyProtectedMember
context = ssl._create_unverified_context() if verify is False else None
conn = connection(req.host, timeout=timeout, context=context)
if req.type == "https":
# noinspection PyProtectedMember
context = ssl._create_unverified_context() if verify is False else None
conn = HTTPSConnection(req.host, timeout=timeout, context=context)
else:
conn = HTTPConnection(req.host, timeout=timeout)

# Make first connection to server
response = self.send_request(conn, req)

# Add connection to the pool if the response is not set to close
Expand Down Expand Up @@ -577,7 +583,7 @@ def send_request(conn, req):

def close(self):
"""Close all persistent connections and remove."""
for _, pool in self.request_handler.values():
for pool in self.request_handler.values():
for key in list(pool.keys()):
conn = pool.pop(key)
conn.close()
Expand Down

0 comments on commit 4bb2275

Please sign in to comment.