Connections are context managers.

Resolves #13.

Author: Lukasa

HISTORY.rst

@@ -9,10 +9,12 @@ X.X.X (XXXX-XX-XX)
 - Implemented logging. (`Issue #12`_)
 - Stopped HTTP/2.0 special headers appearing in the response headers.
   (`Issue #16`_)
+- `HTTP20Connection` objects are now context managers. (`Issue #13`)
 
 
 .. _Issue #12: https://github.com/Lukasa/hyper/issues/12
 .. _Issue #16: https://github.com/Lukasa/hyper/issues/16
+.. _Issue #13: https://github.com/Lukasa/hyper/issues/13
 
 0.0.1 (2014-02-11)
 ++++++++++++++++++

hyper/http20/connection.py

@@ -346,3 +346,13 @@ def _recv_cb(self):
             self.streams[frame.stream_id].receive_frame(frame)
         else:
             self.receive_frame(frame)
+
+
+    # The following two methods are the implementation of the context manager
+    # protocol.
+    def __enter__(self):
+        return self
+
+    def __exit__(self, type, value, tb):
+        self.close()
+        return False  # Never swallow exceptions.

test/test_integration.py

@@ -164,3 +164,36 @@ def socket_handler(listener):
         assert len(frames[-1].data) == 64
 
         self.tear_down()
+
+    def test_connection_context_manager(self):
+        self.set_up()
+
+        data = []
+        send_event = threading.Event()
+
+        def socket_handler(listener):
+            sock = listener.accept()[0]
+
+            # We should get two packets: one connection header string, one
+            # SettingsFrame.
+            first = sock.recv(65535)
+            second = sock.recv(65535)
+            data.append(first)
+            data.append(second)
+
+            # We need to send back a SettingsFrame.
+            f = SettingsFrame(0)
+            sock.send(f.serialize())
+
+            send_event.set()
+            sock.close()
+
+        self._start_server(socket_handler)
+        with HTTP20Connection(self.host, self.port) as conn:
+            conn.connect()
+            send_event.wait()
+
+        # Check that we closed the connection.
+        assert conn._sock == None
+
+        self.tear_down()