1#!/usr/bin/env python3
2# Copyright (C) 2020 The Android Open Source Project
3#
4# Licensed under the Apache License, Version 2.0 (the "License");
5# you may not use this file except in compliance with the License.
6# You may obtain a copy of the License at
7#
8#      http://www.apache.org/licenses/LICENSE-2.0
9#
10# Unless required by applicable law or agreed to in writing, software
11# distributed under the License is distributed on an "AS IS" BASIS,
12# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
13# See the License for the specific language governing permissions and
14# limitations under the License.
15
16import unittest
17
18from trace_processor.api import TraceProcessor, TraceProcessorException
19from trace_processor.protos import ProtoFactory
20
21
22class TestQueryResultIterator(unittest.TestCase):
23  # The numbers input into cells correspond the the CellType enum values
24  # defined under trace_processor.proto
25  CELL_VARINT = ProtoFactory().CellsBatch().CELL_VARINT
26  CELL_STRING = ProtoFactory().CellsBatch().CELL_STRING
27  CELL_INVALID = ProtoFactory().CellsBatch().CELL_INVALID
28
29  def test_one_batch(self):
30    int_values = [100, 200]
31    str_values = ['bar1', 'bar2']
32
33    batch = ProtoFactory().CellsBatch()
34    batch.cells.extend([
35        TestQueryResultIterator.CELL_STRING,
36        TestQueryResultIterator.CELL_VARINT,
37        TestQueryResultIterator.CELL_STRING, TestQueryResultIterator.CELL_VARINT
38    ])
39    batch.varint_cells.extend(int_values)
40    batch.string_cells = "\0".join(str_values) + "\0"
41    batch.is_last_batch = True
42
43    qr_iterator = TraceProcessor.QueryResultIterator(['foo_id', 'foo_num'],
44                                                     [batch])
45
46    for num, row in enumerate(qr_iterator):
47      self.assertEqual(row.foo_id, str_values[num])
48      self.assertEqual(row.foo_num, int_values[num])
49
50  def test_many_batches(self):
51    int_values = [100, 200, 300, 400]
52    str_values = ['bar1', 'bar2', 'bar3', 'bar4']
53
54    batch_1 = ProtoFactory().CellsBatch()
55    batch_1.cells.extend([
56        TestQueryResultIterator.CELL_STRING,
57        TestQueryResultIterator.CELL_VARINT,
58        TestQueryResultIterator.CELL_STRING, TestQueryResultIterator.CELL_VARINT
59    ])
60    batch_1.varint_cells.extend(int_values[:2])
61    batch_1.string_cells = "\0".join(str_values[:2]) + "\0"
62    batch_1.is_last_batch = False
63
64    batch_2 = ProtoFactory().CellsBatch()
65    batch_2.cells.extend([
66        TestQueryResultIterator.CELL_STRING,
67        TestQueryResultIterator.CELL_VARINT,
68        TestQueryResultIterator.CELL_STRING, TestQueryResultIterator.CELL_VARINT
69    ])
70    batch_2.varint_cells.extend(int_values[2:])
71    batch_2.string_cells = "\0".join(str_values[2:]) + "\0"
72    batch_2.is_last_batch = True
73
74    qr_iterator = TraceProcessor.QueryResultIterator(['foo_id', 'foo_num'],
75                                                     [batch_1, batch_2])
76
77    for num, row in enumerate(qr_iterator):
78      self.assertEqual(row.foo_id, str_values[num])
79      self.assertEqual(row.foo_num, int_values[num])
80
81  def test_empty_batch(self):
82    batch = ProtoFactory().CellsBatch()
83    batch.is_last_batch = True
84
85    qr_iterator = TraceProcessor.QueryResultIterator([], [batch])
86
87    for num, row in enumerate(qr_iterator):
88      self.assertIsNone(row.foo_id)
89      self.assertIsNone(row.foo_num)
90
91  def test_invalid_batch(self):
92    batch = ProtoFactory().CellsBatch()
93
94    qr_iterator = TraceProcessor.QueryResultIterator([], [batch])
95
96    # Since the batch isn't defined as the last batch, the QueryResultsIterator
97    # expects another batch and thus raises IndexError as no next batch exists.
98    with self.assertRaises(IndexError):
99      for row in qr_iterator:
100        pass
101
102  def test_incorrect_cells_batch(self):
103    str_values = ['bar1', 'bar2']
104
105    batch = ProtoFactory().CellsBatch()
106    batch.cells.extend([
107        TestQueryResultIterator.CELL_STRING,
108        TestQueryResultIterator.CELL_VARINT,
109        TestQueryResultIterator.CELL_STRING, TestQueryResultIterator.CELL_VARINT
110    ])
111    batch.string_cells = "\0".join(str_values) + "\0"
112    batch.is_last_batch = True
113
114    qr_iterator = TraceProcessor.QueryResultIterator(['foo_id', 'foo_num'],
115                                                     [batch])
116
117    # The batch specifies there ought to be 2 cells of type VARINT and 2 cells
118    # of type STRING, but there are no string cells defined in the batch. Thus
119    # an IndexError occurs as it tries to access the empty string cells list.
120    with self.assertRaises(IndexError):
121      for row in qr_iterator:
122        pass
123
124  def test_incorrect_columns_batch(self):
125    batch = ProtoFactory().CellsBatch()
126    batch.cells.extend([
127        TestQueryResultIterator.CELL_VARINT, TestQueryResultIterator.CELL_VARINT
128    ])
129    batch.varint_cells.extend([100, 200])
130    batch.is_last_batch = True
131
132    qr_iterator = TraceProcessor.QueryResultIterator(
133        ['foo_id', 'foo_num', 'foo_dur', 'foo_ms'], [batch])
134
135    # It's always the case that the number of cells is a multiple of the number
136    # of columns. However, here this is clearly not the case, so when the
137    # iterator tries to access the cell for the third column, it raises an
138    # IndexError due to having exhausted the cells list.
139    with self.assertRaises(IndexError):
140      for row in qr_iterator:
141        pass
142
143  def test_invalid_cell_type(self):
144    batch = ProtoFactory().CellsBatch()
145    batch.cells.extend([
146        TestQueryResultIterator.CELL_INVALID,
147        TestQueryResultIterator.CELL_VARINT
148    ])
149    batch.varint_cells.extend([100, 200])
150    batch.is_last_batch = True
151
152    qr_iterator = TraceProcessor.QueryResultIterator(['foo_id', 'foo_num'],
153                                                     [batch])
154
155    # In this batch we declare the columns types to be CELL_INVALID,
156    # CELL_VARINT but that doesn't match the data which are both ints*
157    # so we should raise a TraceProcessorException.
158    with self.assertRaises(TraceProcessorException):
159      for row in qr_iterator:
160        pass
161
162  def test_one_batch_as_pandas(self):
163    int_values = [100, 200]
164    str_values = ['bar1', 'bar2']
165
166    batch = ProtoFactory().CellsBatch()
167    batch.cells.extend([
168        TestQueryResultIterator.CELL_STRING,
169        TestQueryResultIterator.CELL_VARINT,
170        TestQueryResultIterator.CELL_STRING, TestQueryResultIterator.CELL_VARINT
171    ])
172    batch.varint_cells.extend(int_values)
173    batch.string_cells = "\0".join(str_values) + "\0"
174    batch.is_last_batch = True
175
176    qr_iterator = TraceProcessor.QueryResultIterator(['foo_id', 'foo_num'],
177                                                     [batch])
178
179    qr_df = qr_iterator.as_pandas_dataframe()
180    for num, row in qr_df.iterrows():
181      self.assertEqual(row['foo_id'], str_values[num])
182      self.assertEqual(row['foo_num'], int_values[num])
183
184  def test_many_batches_as_pandas(self):
185    int_values = [100, 200, 300, 400]
186    str_values = ['bar1', 'bar2', 'bar3', 'bar4']
187
188    batch_1 = ProtoFactory().CellsBatch()
189    batch_1.cells.extend([
190        TestQueryResultIterator.CELL_STRING,
191        TestQueryResultIterator.CELL_VARINT,
192        TestQueryResultIterator.CELL_STRING, TestQueryResultIterator.CELL_VARINT
193    ])
194    batch_1.varint_cells.extend(int_values[:2])
195    batch_1.string_cells = "\0".join(str_values[:2]) + "\0"
196    batch_1.is_last_batch = False
197
198    batch_2 = ProtoFactory().CellsBatch()
199    batch_2.cells.extend([
200        TestQueryResultIterator.CELL_STRING,
201        TestQueryResultIterator.CELL_VARINT,
202        TestQueryResultIterator.CELL_STRING, TestQueryResultIterator.CELL_VARINT
203    ])
204    batch_2.varint_cells.extend(int_values[2:])
205    batch_2.string_cells = "\0".join(str_values[2:]) + "\0"
206    batch_2.is_last_batch = True
207
208    qr_iterator = TraceProcessor.QueryResultIterator(['foo_id', 'foo_num'],
209                                                     [batch_1, batch_2])
210
211    qr_df = qr_iterator.as_pandas_dataframe()
212    for num, row in qr_df.iterrows():
213      self.assertEqual(row['foo_id'], str_values[num])
214      self.assertEqual(row['foo_num'], int_values[num])
215
216  def test_empty_batch_as_pandas(self):
217    batch = ProtoFactory().CellsBatch()
218    batch.is_last_batch = True
219
220    qr_iterator = TraceProcessor.QueryResultIterator([], [batch])
221
222    qr_df = qr_iterator.as_pandas_dataframe()
223    for num, row in qr_df.iterrows():
224      self.assertEqual(row['foo_id'], str_values[num])
225      self.assertEqual(row['foo_num'], int_values[num])
226
227  def test_invalid_batch_as_pandas(self):
228    batch = ProtoFactory().CellsBatch()
229
230    qr_iterator = TraceProcessor.QueryResultIterator([], [batch])
231
232    # Since the batch isn't defined as the last batch, the QueryResultsIterator
233    # expects another batch and thus raises IndexError as no next batch exists.
234    with self.assertRaises(IndexError):
235      qr_df = qr_iterator.as_pandas_dataframe()
236
237  def test_incorrect_cells_batch_as_pandas(self):
238    str_values = ['bar1', 'bar2']
239
240    batch = ProtoFactory().CellsBatch()
241    batch.cells.extend([
242        TestQueryResultIterator.CELL_STRING,
243        TestQueryResultIterator.CELL_VARINT,
244        TestQueryResultIterator.CELL_STRING, TestQueryResultIterator.CELL_VARINT
245    ])
246    batch.string_cells = "\0".join(str_values) + "\0"
247    batch.is_last_batch = True
248
249    qr_iterator = TraceProcessor.QueryResultIterator(['foo_id', 'foo_num'],
250                                                     [batch])
251
252    # The batch specifies there ought to be 2 cells of type VARINT and 2 cells
253    # of type STRING, but there are no string cells defined in the batch. Thus
254    # an IndexError occurs as it tries to access the empty string cells list.
255    with self.assertRaises(IndexError):
256      qr_df = qr_iterator.as_pandas_dataframe()
257
258  def test_incorrect_columns_batch_as_pandas(self):
259    batch = ProtoFactory().CellsBatch()
260    batch.cells.extend([
261        TestQueryResultIterator.CELL_VARINT, TestQueryResultIterator.CELL_VARINT
262    ])
263    batch.varint_cells.extend([100, 200])
264    batch.is_last_batch = True
265
266    qr_iterator = TraceProcessor.QueryResultIterator(
267        ['foo_id', 'foo_num', 'foo_dur', 'foo_ms'], [batch])
268
269    # It's always the case that the number of cells is a multiple of the number
270    # of columns. However, here this is clearly not the case, so when the
271    # iterator tries to access the cell for the third column, it raises an
272    # IndexError due to having exhausted the cells list.
273    with self.assertRaises(IndexError):
274      qr_df = qr_iterator.as_pandas_dataframe()
275
276  def test_invalid_cell_type_as_pandas(self):
277    batch = ProtoFactory().CellsBatch()
278    batch.cells.extend([
279        TestQueryResultIterator.CELL_INVALID,
280        TestQueryResultIterator.CELL_VARINT
281    ])
282    batch.varint_cells.extend([100, 200])
283    batch.is_last_batch = True
284
285    qr_iterator = TraceProcessor.QueryResultIterator(['foo_id', 'foo_num'],
286                                                     [batch])
287
288    # In this batch we declare the columns types to be CELL_INVALID,
289    # CELL_VARINT but that doesn't match the data which are both ints*
290    # so we should raise a TraceProcessorException.
291    with self.assertRaises(TraceProcessorException):
292      qr_df = qr_iterator.as_pandas_dataframe()
293