Merge pull request #185 from lelutin/enable_wth_statement

Enable with statement

README.rst

@@ -97,6 +97,23 @@ The first thing you need to do is to connect with PuppetDB:
    >>> from pypuppetdb import connect
    >>> db = connect()
 
+Once you're all done with requests to PuppetDB, you can explicitly close all
+HTTP connections. This can be useful if you made those connections through a
+tunnel, which might be tracking open connections:
+
+.. code-block:: python
+
+   >>> db.disconnect()
+
+You can also use the `with` statement to enclose the work done on PuppetDB.
+This will ensure that all HTTP connections are closed explicitly when we're
+done:
+
+.. code-block:: python
+
+   >>> with connect() as db:
+   >>>   # ..
+
 Nodes
 -----
 

docs/quickstart.rst

@@ -15,6 +15,23 @@ The first thing you need to do is to connect with PuppetDB:
    >>> from pypuppetdb import connect
    >>> db = connect()
 
+Once you're all done with requests to PuppetDB, you can explicitly close all
+HTTP connections. This can be useful if you made those connections through a
+tunnel, which might be tracking open connections:
+
+.. code-block:: python
+
+   >>> db.disconnect()
+
+You can also use the `with` statement to enclose the work done on PuppetDB.
+This will ensure that all HTTP connections are closed explicitly when we're
+done:
+
+.. code-block:: python
+
+   >>> with connect() as db:
+   >>>   # ..
+
 Nodes
 -----
 

pypuppetdb/api.py

@@ -194,6 +194,23 @@ def __init__(self, host='localhost', port=8080, ssl_verify=True,
         else:
             self.protocol = 'http'
 
+    def disconnect(self):
+        """Close all connections that this class opened up."""
+        # If we don't explicitly close connections, we might cause other
+        # functions or libraries to hang on the open connections. This happens
+        # for example with using paramiko to tunnel PuppetDB connections
+        # through ssh.
+        self._session.close()
+
+    def __enter__(self):
+        """Set up environment for 'with' statement."""
+        # Once this class has been instantiated, there's nothing more required
+        return self
+
+    def __exit__(self, type, value, trace):
+        """Tear down connections."""
+        self.disconnect()
+
     @property
     def version(self):
         """The version of the API we're querying against.

tests/test_connect.py

@@ -1,6 +1,20 @@
+import mock
+
 import pypuppetdb
 
 
-def test_connect_api():
+@mock.patch('pypuppetdb.api.requests.Session.close')
+def test_connect_api(session_close):
     puppetdb = pypuppetdb.connect()
     assert puppetdb.version == 'v4'
+
+    puppetdb.disconnect()
+    assert session_close.called
+
+
+@mock.patch('pypuppetdb.api.requests.Session.close')
+def test_connect_with_statement(session_close):
+    with pypuppetdb.connect() as puppetdb:
+        assert puppetdb.version == 'v4'
+
+    assert session_close.called