2 files changed, 59 insertions, 46 deletions
diff --git a/src/leap/common/http.py b/src/leap/common/http.py
index 56938b4..f67507d 100644
--- a/src/leap/common/http.py
+++ b/src/leap/common/http.py
@@ -56,6 +56,50 @@ __all__ = ["HTTPClient"]
 DEFAULT_HTTP_TIMEOUT = 30  # seconds
 
 
+class _HTTP11ClientFactory(HTTP11ClientFactory):
+    """
+    A timeout-able HTTP 1.1 client protocol factory.
+    """
+
+    def __init__(self, quiescentCallback, timeout):
+        """
+        :param quiescentCallback: The quiescent callback to be passed to
+                                  protocol instances, used to return them to
+                                  the connection pool.
+        :type quiescentCallback: callable(Protocol)
+        :param timeout: The timeout, in seconds, for requests made by
+                        protocols created by this factory.
+        :type timeout: float
+        """
+        HTTP11ClientFactory.__init__(self, quiescentCallback)
+        self._timeout = timeout
+
+    def buildProtocol(self, _):
+        """
+        Build the HTTP 1.1 client protocol.
+        """
+        return _HTTP11ClientProtocol(self._quiescentCallback, self._timeout)
+
+
+class _HTTPConnectionPool(HTTPConnectionPool):
+    """
+    A timeout-able HTTP connection pool.
+    """
+
+    _factory = _HTTP11ClientFactory
+
+    def __init__(self, reactor, persistent, timeout, maxPersistentPerHost=10):
+        HTTPConnectionPool.__init__(self, reactor, persistent=persistent)
+        self.maxPersistentPerHost = maxPersistentPerHost
+        self._timeout = timeout
+
+    def _newConnection(self, key, endpoint):
+        def quiescentCallback(protocol):
+            self._putConnection(key, protocol)
+        factory = self._factory(quiescentCallback, timeout=self._timeout)
+        return endpoint.connect(factory)
+
+
 class HTTPClient(object):
     """
     HTTP client done the twisted way, with a main focus on pinning the SSL
@@ -69,7 +113,14 @@ class HTTPClient(object):
     in order to avoid resource abuse on huge requests batches.
     """
 
-    def __init__(self, cert_file=None, timeout=DEFAULT_HTTP_TIMEOUT):
+    _pool = _HTTPConnectionPool(
+        reactor,
+        persistent=True,
+        timeout = DEFAULT_HTTP_TIMEOUT,
+        maxPersistentPerHost=10
+    )
+
+    def __init__(self, cert_file=None, timeout=DEFAULT_HTTP_TIMEOUT, pool=None):
         """
         Init the HTTP client
 
@@ -84,7 +135,7 @@ class HTTPClient(object):
         """
 
         self._timeout = timeout
-        self._pool = self._createPool()
+        self._pool = pool if pool is not None else self._pool
         self._agent = Agent(
             reactor,
             get_compatible_ssl_context_factory(cert_file),
@@ -278,46 +329,3 @@ class _HTTP11ClientProtocol(HTTP11ClientProtocol):
         HTTP11ClientProtocol.dataReceived(self, bytes)
         if self._timeoutCall and self._timeoutCall.active():
             self._timeoutCall.reset(self._timeout)
-
-
-class _HTTP11ClientFactory(HTTP11ClientFactory):
-    """
-    A timeout-able HTTP 1.1 client protocol factory.
-    """
-
-    def __init__(self, quiescentCallback, timeout):
-        """
-        :param quiescentCallback: The quiescent callback to be passed to
-                                  protocol instances, used to return them to
-                                  the connection pool.
-        :type quiescentCallback: callable(Protocol)
-        :param timeout: The timeout, in seconds, for requests made by
-                        protocols created by this factory.
-        :type timeout: float
-        """
-        HTTP11ClientFactory.__init__(self, quiescentCallback)
-        self._timeout = timeout
-
-    def buildProtocol(self, _):
-        """
-        Build the HTTP 1.1 client protocol.
-        """
-        return _HTTP11ClientProtocol(self._quiescentCallback, self._timeout)
-
-
-class _HTTPConnectionPool(HTTPConnectionPool):
-    """
-    A timeout-able HTTP connection pool.
-    """
-
-    _factory = _HTTP11ClientFactory
-
-    def __init__(self, reactor, persistent, timeout):
-        HTTPConnectionPool.__init__(self, reactor, persistent=persistent)
-        self._timeout = timeout
-
-    def _newConnection(self, key, endpoint):
-        def quiescentCallback(protocol):
-            self._putConnection(key, protocol)
-        factory = self._factory(quiescentCallback, timeout=self._timeout)
-        return endpoint.connect(factory)
diff --git a/src/leap/common/tests/test_http.py b/src/leap/common/tests/test_http.py
index e240ca3..a586fd1 100644
--- a/src/leap/common/tests/test_http.py
+++ b/src/leap/common/tests/test_http.py
@@ -47,7 +47,12 @@ class HTTPClientTest(BaseLeapTest):
         self.assertEquals(client._agent._pool, client2._agent._pool, "Pool was not reused by default")
 
     def test_agent_can_have_dedicated_custom_pool(self):
-        custom_pool = http.createPool(maxPersistentPerHost=42, persistent=False)
+        custom_pool = http._HTTPConnectionPool(
+            None,
+            timeout=10,
+            maxPersistentPerHost=42,
+            persistent=False
+        )
         self.assertEquals(custom_pool.maxPersistentPerHost, 42,
                           "Custom persistent connections limit is not being respected")
         self.assertFalse(custom_pool.persistent,