From f33706ae172c709fb2a20a7e7534cfda5150bf70 Mon Sep 17 00:00:00 2001
From: davebshow <davebshow@gmail.com>
Date: Wed, 20 May 2015 13:20:57 -0400
Subject: [PATCH] fixing up client conn handling

---
 README.md                    |  2 +-
 aiogremlin/__init__.py       |  2 +-
 aiogremlin/client.py         |  8 +++-----
 aiogremlin/contextmanager.py |  2 ++
 changes.txt                  |  2 ++
 setup.py                     |  2 +-
 tests/tests.py               | 22 ++++++++++++++++++++--
 7 files changed, 30 insertions(+), 10 deletions(-)

diff --git a/README.md b/README.md
index 9da6f42..868847c 100644
--- a/README.md
+++ b/README.md
@@ -1,4 +1,4 @@
-# aiogremlin 0.0.6 [(gizmo grew up)](https://pypi.python.org/pypi/gizmo/0.1.12)
+# aiogremlin 0.0.8 [(gizmo grew up)](https://pypi.python.org/pypi/gizmo/0.1.12)
 
 `aiogremlin` is a **Python 3** driver for the the [Tinkerpop 3 Gremlin Server](http://www.tinkerpop.com/docs/3.0.0.M7/#gremlin-server). This module is built on [Asyncio](https://docs.python.org/3/library/asyncio.html) and [aiohttp](http://aiohttp.readthedocs.org/en/v0.15.3/index.html) `aiogremlin` is currently in **alpha** mode, but all major functionality has test coverage.
 
diff --git a/aiogremlin/__init__.py b/aiogremlin/__init__.py
index 3f6fa2f..0fc5fac 100644
--- a/aiogremlin/__init__.py
+++ b/aiogremlin/__init__.py
@@ -3,4 +3,4 @@ from .client import *
 from .exceptions import *
 from .pool import *
 from .protocol import *
-__version__ = "0.0.7"
+__version__ = "0.0.8"
diff --git a/aiogremlin/client.py b/aiogremlin/client.py
index e970826..224e193 100644
--- a/aiogremlin/client.py
+++ b/aiogremlin/client.py
@@ -71,15 +71,13 @@ class GremlinClient:
         self._connector = connector
         self._factory = factory or GremlinFactory(connector=self._connector)
         self._conn = connection
-        if self._pool is None and self._conn is None:
-            self._connected = False
-            self._conn = asyncio.async(self._connect(), loop=self._loop)
-        elif pool is not None:
+        if pool is not None:
             self._connected = self._pool._connected
-        elif self._conn is not None and not connection._closed:
+        elif self._conn and not getattr(connection, "closed", True):
             self._connected = True
         else:
             self._connected = False
+            self._conn = asyncio.async(self._connect(), loop=self._loop)
         if verbose:
             logger.setLevel(INFO)
 
diff --git a/aiogremlin/contextmanager.py b/aiogremlin/contextmanager.py
index b5ecade..ff02537 100644
--- a/aiogremlin/contextmanager.py
+++ b/aiogremlin/contextmanager.py
@@ -6,6 +6,8 @@ class ConnectionContextManager:
         self._conn = conn
 
     def __enter__(self):
+        if self._conn.closed:
+            raise RuntimeError("Connection closed unexpectedly.")
         return self._conn
 
     def __exit__(self, exception_type, exception_value, traceback):
diff --git a/changes.txt b/changes.txt
index 20e3974..d7b4cba 100644
--- a/changes.txt
+++ b/changes.txt
@@ -5,3 +5,5 @@
 0.0.5 - 5/13/2015: Using EofStream terminator technique to prepare for 3.0.0.M9
 0.0.7 - 5/20/2015: Full integration with aiohttp. Added Session object for HTTP
                    connection pooling. Added a context manager.
+0.0.8 - 5/20/2015: Fixed bug in client constructor logic. Tested different client
+                   connection handling techniques.
diff --git a/setup.py b/setup.py
index 8e12c39..44e9048 100644
--- a/setup.py
+++ b/setup.py
@@ -3,7 +3,7 @@ from setuptools import setup
 
 setup(
     name="aiogremlin",
-    version="0.0.7",
+    version="0.0.8",
     url="",
     license="MIT",
     author="davebshow",
diff --git a/tests/tests.py b/tests/tests.py
index 4fb908f..7937833 100644
--- a/tests/tests.py
+++ b/tests/tests.py
@@ -275,13 +275,31 @@ class ContextMngrTest(unittest.TestCase):
         def go():
             with (yield from self.pool) as conn:
                 gc = GremlinClient(connection=conn, loop=self.loop)
-                resp = yield from gc.execute("1 + 1")
-                self.assertEqual(resp[0].data[0], 2)
+                resp = yield from gc.submit("1 + 1")
+                self.assertEqual(conn, resp.stream._conn)
+                result = yield from resp.get()
+                self.assertEqual(result[0].data[0], 2)
+
                 self.pool.release(conn)
             # Test that connection was closed
             yield from self._check_closed()
         self.loop.run_until_complete(go())
 
+    def test_connection_manager_with_client_closed_conn(self):
+        @asyncio.coroutine
+        def go():
+            with (yield from self.pool) as conn:
+                conn._closing = True
+                conn._close()
+                gc = GremlinClient(connection=conn, loop=self.loop)
+                resp = yield from gc.submit("1 + 1")
+                self.assertNotEqual(conn, resp.stream._conn)
+                result = yield from resp.get()
+                self.assertEqual(result[0].data[0], 2)
+                yield from resp.stream._conn.close()
+            # Test that connection was closed
+        self.loop.run_until_complete(go())
+
     def test_connection_manager_error(self):
         results = []
         @asyncio.coroutine
-- 
GitLab