From 2045e89f24535d846d482ac574748d88285090dc Mon Sep 17 00:00:00 2001
From: Jim Paris <jim@jtan.com>
Date: Sat, 16 Feb 2013 18:47:07 -0500
Subject: [PATCH] client: Add context manager functionality, test closing

---
 nilmdb/client/client.py     |  8 ++++++++
 nilmdb/client/httpclient.py |  4 ++--
 tests/test_client.py        | 36 +++++++++++++++++++++++++++++++-----
 3 files changed, 41 insertions(+), 7 deletions(-)

diff --git a/nilmdb/client/client.py b/nilmdb/client/client.py
index bba733f..3e1d4e5 100644
--- a/nilmdb/client/client.py
+++ b/nilmdb/client/client.py
@@ -19,11 +19,19 @@ class Client(object):
     def __init__(self, url):
         self.http = nilmdb.client.httpclient.HTTPClient(url)
 
+    # __enter__/__exit__ allow this class to be a context manager
+    def __enter__(self):
+        return self
+
+    def __exit__(self, exc_type, exc_value, traceback):
+        self.close()
+
     def _json_param(self, data):
         """Return compact json-encoded version of parameter"""
         return json.dumps(data, separators=(',',':'))
 
     def close(self):
+        """Close the connection; safe to call multiple times"""
         self.http.close()
 
     def geturl(self):
diff --git a/nilmdb/client/httpclient.py b/nilmdb/client/httpclient.py
index f8530f5..dee8ab4 100644
--- a/nilmdb/client/httpclient.py
+++ b/nilmdb/client/httpclient.py
@@ -80,11 +80,11 @@ class HTTPClient(object):
                 self._status = int(data.split(" ")[1])
             self._headers += data
         self.curl.setopt(pycurl.HEADERFUNCTION, header_callback)
-        def func(callback):
+        def perform(callback):
             self.curl.setopt(pycurl.WRITEFUNCTION, callback)
             self.curl.perform()
         try:
-            with nilmdb.utils.Iteratorizer(func, curl_hack = True) as it:
+            with nilmdb.utils.Iteratorizer(perform, curl_hack = True) as it:
                 for i in it:
                     if self._status == 200:
                         # If we had a 200 response, yield the data to caller.
diff --git a/tests/test_client.py b/tests/test_client.py
index d1593b0..838abc6 100644
--- a/tests/test_client.py
+++ b/tests/test_client.py
@@ -18,6 +18,7 @@ import simplejson as json
 import unittest
 import warnings
 import resource
+import time
 
 from testutil.helpers import *
 
@@ -101,8 +102,10 @@ class TestClient(object):
                                     ["/newton/zzz/rawnotch", "RawNotchedData"]
                                     ])
         # Match just one type or one path
-        eq_(client.stream_list(layout="RawData"), [ ["/newton/raw", "RawData"] ])
-        eq_(client.stream_list(path="/newton/raw"), [ ["/newton/raw", "RawData"] ])
+        eq_(client.stream_list(layout="RawData"),
+            [ ["/newton/raw", "RawData"] ])
+        eq_(client.stream_list(path="/newton/raw"),
+            [ ["/newton/raw", "RawData"] ])
 
         # Try messing with resource limits to trigger errors and get
         # more coverage.  Here, make it so we can only create files 1
@@ -131,9 +134,10 @@ class TestClient(object):
         client.stream_update_metadata("/newton/raw", meta3)
         eq_(client.stream_get_metadata("/newton/prep"), meta1)
         eq_(client.stream_get_metadata("/newton/raw"), meta1)
-        eq_(client.stream_get_metadata("/newton/raw", [ "description" ] ), meta2)
-        eq_(client.stream_get_metadata("/newton/raw", [ "description",
-                                                        "v_scale" ] ), meta1)
+        eq_(client.stream_get_metadata("/newton/raw",
+                                       [ "description" ] ), meta2)
+        eq_(client.stream_get_metadata("/newton/raw",
+                                       [ "description", "v_scale" ] ), meta1)
 
         # missing key
         eq_(client.stream_get_metadata("/newton/raw", "descr"),
@@ -356,3 +360,25 @@ class TestClient(object):
         eq_(client.stream_get_metadata(raw[0]), meta1)
         eq_(client.stream_get_metadata(raw[0], [ "alpha" ]), meta2)
         eq_(client.stream_get_metadata(raw[0], [ "alpha", "β" ]), meta1)
+
+    def test_client_9_closing(self):
+        # Make sure we actually close sockets correctly.  New
+        # connections will block for a while if they're not, since the
+        # server will stop accepting new connections.
+        for test in [1, 2]:
+            start = time.time()
+            for i in range(50):
+                if time.time() - start > 15:
+                    raise AssertionError("Connections seem to be blocking... "
+                                         "probably not closing properly.")
+                if test == 1:
+                    # explicit close
+                    client = nilmdb.Client(url = "http://localhost:12380/")
+                    with assert_raises(ClientError) as e:
+                        client.stream_remove("/newton/prep", 123, 120)
+                    client.close() # remove this to see the failure
+                elif test == 2:
+                    # use the context manager
+                    with nilmdb.Client(url = "http://localhost:12380/") as c:
+                        with assert_raises(ClientError) as e:
+                            c.stream_remove("/newton/prep", 123, 120)