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