Result consumption tests

Author: technige

neo4j/v1/session.py

@@ -206,12 +206,12 @@ def on_failure(metadata):
         pull_all_response.on_failure = on_failure
 
     def __iter__(self):
-        while self._buffer and not self._consumed:
+        while self._buffer:
             values = self._buffer.popleft()
             yield Record(self.keys(), tuple(map(hydrated, values)))
         while not self._consumed:
             self.connection.fetch()
-            if self._buffer:
+            while self._buffer:
                 values = self._buffer.popleft()
                 yield Record(self.keys(), tuple(map(hydrated, values)))
 
@@ -223,14 +223,18 @@ def keys(self):
             self.connection.fetch()
         return self._keys
 
+    def buffer(self):
+        if self.connection and not self.connection.closed:
+            while not self._consumed:
+                self.connection.fetch()
+            self.connection = None
+
     def consume(self):
         """ Consume the remainder of this result and return the
         summary.
         """
         if self.connection and not self.connection.closed:
-            fetch = self.connection.fetch
-            while not self._consumed:
-                fetch()
+            list(self)
             self.connection = None
         return self._summary
 
@@ -475,7 +479,7 @@ def close(self):
         """ Recycle this session through the driver it came from.
         """
         if self.last_result:
-            self.last_result.consume()
+            self.last_result.buffer()
         self.driver.recycle(self)
 
     def begin_transaction(self):

test/test_session.py

@@ -501,3 +501,67 @@ def test_can_rollback_transaction_using_with_block(self):
             result = session.run("MATCH (a) WHERE id(a) = {n} "
                                  "RETURN a.foo", {"n": node_id})
             assert len(list(result)) == 0
+
+
+class ResultConsumptionTestCase(ServerTestCase):
+
+    def setUp(self):
+        self.driver = GraphDatabase.driver("bolt://localhost", auth=auth_token)
+
+    def test_can_consume_result_immediately(self):
+        session = self.driver.session()
+        tx = session.begin_transaction()
+        result = tx.run("UNWIND range(1, 3) AS n RETURN n")
+        assert [record[0] for record in result] == [1, 2, 3]
+        tx.commit()
+        session.close()
+
+    def test_can_consume_result_after_commit(self):
+        session = self.driver.session()
+        tx = session.begin_transaction()
+        result = tx.run("UNWIND range(1, 3) AS n RETURN n")
+        tx.commit()
+        assert [record[0] for record in result] == [1, 2, 3]
+        session.close()
+
+    def test_can_consume_result_after_rollback(self):
+        session = self.driver.session()
+        tx = session.begin_transaction()
+        result = tx.run("UNWIND range(1, 3) AS n RETURN n")
+        tx.rollback()
+        assert [record[0] for record in result] == [1, 2, 3]
+        session.close()
+
+    def test_can_consume_result_after_session_close(self):
+        session = self.driver.session()
+        tx = session.begin_transaction()
+        result = tx.run("UNWIND range(1, 3) AS n RETURN n")
+        tx.commit()
+        session.close()
+        assert [record[0] for record in result] == [1, 2, 3]
+
+    def test_can_consume_result_after_session_reuse(self):
+        session = self.driver.session()
+        tx = session.begin_transaction()
+        result_a = tx.run("UNWIND range(1, 3) AS n RETURN n")
+        tx.commit()
+        session.close()
+        session = self.driver.session()
+        tx = session.begin_transaction()
+        result_b = tx.run("UNWIND range(4, 6) AS n RETURN n")
+        tx.commit()
+        session.close()
+        assert [record[0] for record in result_a] == [1, 2, 3]
+        assert [record[0] for record in result_b] == [4, 5, 6]
+
+    def test_can_consume_result_after_session_with_error(self):
+        session = self.driver.session()
+        with self.assertRaises(CypherError):
+            session.run("X").consume()
+        session.close()
+        session = self.driver.session()
+        tx = session.begin_transaction()
+        result = tx.run("UNWIND range(1, 3) AS n RETURN n")
+        tx.commit()
+        session.close()
+        assert [record[0] for record in result] == [1, 2, 3]