1  import io
 2  import json
 3  from itertools import zip_longest
 4  from unittest import TestCase
 5  from unittest.mock import Mock
 6  
 7  from json_stream.requests import IterableStream, load
 8  
 9  
10  class TestIterableStream(TestCase):
11      def test_read(self):
12          # create some chunks of binary data
13          data = (
14              b"a" * io.DEFAULT_BUFFER_SIZE,
15              b"b" * (io.DEFAULT_BUFFER_SIZE + 1),
16              b"c" * (io.DEFAULT_BUFFER_SIZE - 1),
17          )
18  
19          # stream it and check the result
20          stream = IterableStream(data)
21          self.assertEqual(stream.read(), b"".join(data))
22  
23  
24  class TestLoad(TestCase):
25      @staticmethod
26      def grouper(iterable, n):
27          "Collect data into fixed-length chunks or blocks"
28          # grouper('ABCDEFG', 3, 'x') --> ABC DEF Gxx"
29          args = [iter(iterable)] * n
30          return zip_longest(*args, fillvalue="")
31  
32      def test_load_persistent(self):
33          # requests iter_content returns an iterable of bytes
34          response = Mock()
35          data = json.dumps({
36              "a": "a" * io.DEFAULT_BUFFER_SIZE,
37              "b": "b",
38          })
39          content = ("".join(chunk).encode() for chunk in self.grouper(data, 1024))
40          response.iter_content.return_value = content
41  
42          # load in persistent mode
43          data = load(response, persistent=True)
44  
45          # is streaming
46          self.assertTrue(data.streaming)
47  
48          # check the data
49          self.assertTupleEqual(
50              (
51                  ("a", "a" * io.DEFAULT_BUFFER_SIZE),
52                  ("b", "b"),
53              ),
54              tuple(sorted(data.items())),
55          )
56  
57          # all done?
58          self.assertFalse(data.streaming)
59  
60      def test_load_transient(self):
61          # requests iter_content returns an iterable of bytes
62          response = Mock()
63          data = json.dumps({
64              "a": "a" * io.DEFAULT_BUFFER_SIZE,
65              "b": "b",
66          })
67          content = ("".join(chunk).encode() for chunk in self.grouper(data, 1024))
68          response.iter_content.return_value = content
69  
70          # load in transient mode
71          data = load(response, persistent=False)
72  
73          # is streaming
74          self.assertTrue(data.streaming)
75  
76          # check the data
77          self.assertTupleEqual(
78              (
79                  ("a", "a" * io.DEFAULT_BUFFER_SIZE),
80                  ("b", "b"),
81              ),
82              tuple(sorted(data.items())),
83          )
84  
85          # all done?
86          self.assertFalse(data.streaming)