Merge pull request influxdata#84 from influxdb/fix_21_proposition

aviau · aviau · commit 098892342527 · 2014-11-10T23:53:03.000-05:00
Fix 21 proposition
diff --git a/influxdb/chunked_json.py b/influxdb/chunked_json.py
@@ -0,0 +1,21 @@
+# -*- coding: utf-8 -*-
+
+#
+# Author: Adrian Sampson <adrian@radbox.org>
+# Source: https://gist.github.com/sampsyo/920215
+#
+
+import json
+
+_decoder = json.JSONDecoder()
+
+
+def loads(s):
+    """A generator reading a sequence of JSON values from a string."""
+    while s:
+        s = s.strip()
+        obj, pos = _decoder.raw_decode(s)
+        if not pos:
+            raise ValueError('no JSON object found at %i' % pos)
+        yield obj
+        s = s[pos:]
diff --git a/influxdb/client.py b/influxdb/client.py
@@ -6,6 +6,8 @@
 import socket
 import requests
 
+from influxdb import chunked_json
+
 try:
     xrange
 except NameError:
@@ -322,7 +324,10 @@ def query(self, query, time_precision='s', chunked=False):
             expected_response_code=200
         )
 
-        return response.json()
+        if chunked:
+            return list(chunked_json.loads(response.content.decode()))
+        else:
+            return response.json()
 
     # Creating and Dropping Databases
     #
diff --git a/tests/influxdb/client_test.py b/tests/influxdb/client_test.py
@@ -242,7 +242,7 @@ def test_query(self):
 
     def test_query_chunked(self):
         cli = InfluxDBClient(database='db')
-        example_response = {
+        example_object = {
             'points': [
                 [1415206250119, 40001, 667],
                 [1415206244555, 30001, 7],
@@ -257,17 +257,19 @@ def test_query_chunked(self):
                 'val'
             ]
         }
+        example_response = \
+            json.dumps(example_object) + json.dumps(example_object)
 
         with requests_mock.Mocker() as m:
             m.register_uri(
                 requests_mock.GET,
                 "http://localhost:8086/db/db/series",
-                text=json.dumps(example_response)
+                text=example_response
             )
 
-            self.assertDictEqual(
+            self.assertListEqual(
                 cli.query('select * from foo', chunked=True),
-                example_response
+                [example_object, example_object]
             )
 
     @raises(Exception)
